Bugfix rpc handshake

wrappers/blockcipherrwc.go

@@ -25,6 +25,7 @@ func (c *BlockCipherReadWriteCloserWithAddr) Read(p []byte) (int, error) {
 				return errors.Annotate(err, "Cannot read from socket")
 			}
 			c.Buffer.Write(p[:n])
+			bufferLength = c.Buffer.Len()
 		}
 		c.decryptor.CryptBlocks(c.Buffer.Bytes(), c.Buffer.Bytes())
 
@@ -47,8 +48,12 @@ func (c *BlockCipherReadWriteCloserWithAddr) Close() error {
 	return c.conn.Close()
 }
 
-func (c *BlockCipherReadWriteCloserWithAddr) Addr() *net.TCPAddr {
-	return c.conn.Addr()
+func (c *BlockCipherReadWriteCloserWithAddr) LocalAddr() *net.TCPAddr {
+	return c.conn.LocalAddr()
+}
+
+func (c *BlockCipherReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
+	return c.conn.RemoteAddr()
 }
 
 func NewBlockCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.BlockMode) ReadWriteCloserWithAddr {

wrappers/ctxrwc.go

@@ -48,8 +48,12 @@ func (c *CtxReadWriteCloserWithAddr) Close() error {
 	return c.conn.Close()
 }
 
-func (c *CtxReadWriteCloserWithAddr) Addr() *net.TCPAddr {
-	return c.conn.Addr()
+func (c *CtxReadWriteCloserWithAddr) LocalAddr() *net.TCPAddr {
+	return c.conn.LocalAddr()
+}
+
+func (c *CtxReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
+	return c.conn.RemoteAddr()
 }
 
 // NewCtxRWC returns ReadWriteCloser which respects given context,

wrappers/logrwc.go

@@ -36,8 +36,12 @@ func (l *LogReadWriteCloserWithAddr) Close() error {
 	return err
 }
 
-func (l *LogReadWriteCloserWithAddr) Addr() *net.TCPAddr {
-	return l.conn.Addr()
+func (l *LogReadWriteCloserWithAddr) LocalAddr() *net.TCPAddr {
+	return l.conn.LocalAddr()
+}
+
+func (l *LogReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
+	return l.conn.RemoteAddr()
 }
 
 // NewLogRWC wraps ReadWriteCloser with logger calls.

wrappers/rwcaddr.go

@@ -8,5 +8,6 @@ import (
 type ReadWriteCloserWithAddr interface {
 	io.ReadWriteCloser
 
-	Addr() *net.TCPAddr
+	LocalAddr() *net.TCPAddr
+	RemoteAddr() *net.TCPAddr
 }

wrappers/streamcipherrwc.go

@@ -43,8 +43,12 @@ func (c *StreamCipherReadWriteCloserWithAddr) Close() error {
 	return c.conn.Close()
 }
 
-func (c *StreamCipherReadWriteCloserWithAddr) Addr() *net.TCPAddr {
-	return c.conn.Addr()
+func (c *StreamCipherReadWriteCloserWithAddr) LocalAddr() *net.TCPAddr {
+	return c.conn.LocalAddr()
+}
+
+func (c *StreamCipherReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
+	return c.conn.RemoteAddr()
 }
 
 // NewStreamCipherRWC returns wrapper which transparently

wrappers/timeoutrwc.go

@@ -8,7 +8,9 @@ import (
 )
 
 type TimeoutReadWriteCloserWithAddr struct {
-	conn net.Conn
+	conn       net.Conn
+	publicIPv4 net.IP
+	publicIPv6 net.IP
 }
 
 func (t *TimeoutReadWriteCloserWithAddr) Read(p []byte) (int, error) {
@@ -25,10 +27,29 @@ func (t *TimeoutReadWriteCloserWithAddr) Close() error {
 	return t.conn.Close()
 }
 
-func (t *TimeoutReadWriteCloserWithAddr) Addr() *net.TCPAddr {
+func (t *TimeoutReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
 	return t.conn.RemoteAddr().(*net.TCPAddr)
 }
 
-func NewTimeoutRWC(conn net.Conn) ReadWriteCloserWithAddr {
-	return &TimeoutReadWriteCloserWithAddr{conn}
+func (t *TimeoutReadWriteCloserWithAddr) LocalAddr() *net.TCPAddr {
+	addr := t.conn.LocalAddr().(*net.TCPAddr)
+	newAddr := *addr
+
+	if t.RemoteAddr().IP.To4() != nil {
+		if t.publicIPv4 != nil {
+			newAddr.IP = t.publicIPv4
+		}
+	} else if t.publicIPv6 != nil {
+		newAddr.IP = t.publicIPv6
+	}
+
+	return &newAddr
+}
+
+func NewTimeoutRWC(conn net.Conn, ipv4, ipv6 net.IP) ReadWriteCloserWithAddr {
+	return &TimeoutReadWriteCloserWithAddr{
+		conn:       conn,
+		publicIPv4: ipv4,
+		publicIPv6: ipv6,
+	}
 }

wrappers/trafficrwc.go

@@ -29,8 +29,12 @@ func (t *TrafficReadWriteCloserWithAddr) Close() error {
 	return t.conn.Close()
 }
 
-func (t *TrafficReadWriteCloserWithAddr) Addr() *net.TCPAddr {
-	return t.conn.Addr()
+func (t *TrafficReadWriteCloserWithAddr) LocalAddr() *net.TCPAddr {
+	return t.conn.LocalAddr()
+}
+
+func (t *TrafficReadWriteCloserWithAddr) RemoteAddr() *net.TCPAddr {
+	return t.conn.RemoteAddr()
 }
 
 // NewTrafficRWC wraps ReadWriteCloser to have read/write callbacks.