From 77b186ae6eb7f61ff4e51494c238388b88f68194 Mon Sep 17 00:00:00 2001 From: David Fifield Date: Sat, 24 Sep 2022 12:39:31 -0600 Subject: [PATCH] Have SnowflakeClientConn implement io.WriterTo. By forwarding the method to the inner smux.Stream. This is to prevent io.Copy in the top-level proxy function from allocating a buffer per client. The smux.Stream WriteTo method returns io.EOF on success, contrary to the contract of io.Copy that says it should return nil. Ignore io.EOF in the proxy loop to avoid a log message. /anti-censorship/pluggable-transports/snowflake/-/issues/40177 --- server/lib/snowflake.go | 32 ++++++++++++++++++++++++++------ server/server.go | 2 +- 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/server/lib/snowflake.go b/server/lib/snowflake.go index 469ed91..5f995ac 100644 --- a/server/lib/snowflake.go +++ b/server/lib/snowflake.go @@ -232,7 +232,7 @@ func (l *SnowflakeListener) acceptStreams(conn *kcp.UDPSession) error { } return err } - l.queueConn(&SnowflakeClientConn{Conn: stream, address: addr}) + l.queueConn(&SnowflakeClientConn{stream: stream, address: addr}) } } @@ -280,15 +280,35 @@ func (l *SnowflakeListener) queueConn(conn net.Conn) error { } } -// SnowflakeClientConn is a wrapper for the underlying turbotunnel -// conn. We need to reference our client address map to determine the -// remote address +// SnowflakeClientConn is a wrapper for the underlying turbotunnel conn +// (smux.Stream). It implements the net.Conn and io.WriterTo interfaces. The +// RemoteAddr method is overridden to refer to a real IP address, looked up from +// the client address map, rather than an abstract client ID. type SnowflakeClientConn struct { - net.Conn + stream *smux.Stream address net.Addr } -// RemoteAddr returns the mapped client address of the Snowflake connection +// Forward net.Conn methods, other than RemoteAddr, to the inner stream. +func (conn *SnowflakeClientConn) Read(b []byte) (int, error) { return conn.stream.Read(b) } +func (conn *SnowflakeClientConn) Write(b []byte) (int, error) { return conn.stream.Write(b) } +func (conn *SnowflakeClientConn) Close() error { return conn.stream.Close() } +func (conn *SnowflakeClientConn) LocalAddr() net.Addr { return conn.stream.LocalAddr() } +func (conn *SnowflakeClientConn) SetDeadline(t time.Time) error { return conn.stream.SetDeadline(t) } +func (conn *SnowflakeClientConn) SetReadDeadline(t time.Time) error { + return conn.stream.SetReadDeadline(t) +} +func (conn *SnowflakeClientConn) SetWriteDeadline(t time.Time) error { + return conn.stream.SetWriteDeadline(t) +} + +// RemoteAddr returns the mapped client address of the Snowflake connection. func (conn *SnowflakeClientConn) RemoteAddr() net.Addr { return conn.address } + +// WriteTo implements the io.WriterTo interface by passing the call to the +// underlying smux.Stream. +func (conn *SnowflakeClientConn) WriteTo(w io.Writer) (int64, error) { + return conn.stream.WriteTo(w) +} diff --git a/server/server.go b/server/server.go index ed1e876..484e37a 100644 --- a/server/server.go +++ b/server/server.go @@ -57,7 +57,7 @@ func proxy(local *net.TCPConn, conn net.Conn) { wg.Done() }() go func() { - if _, err := io.Copy(local, conn); err != nil && !errors.Is(err, io.ErrClosedPipe) { + if _, err := io.Copy(local, conn); err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrClosedPipe) { log.Printf("error copying WebSocket to ORPort %v", err) } local.CloseWrite()