feat: add a GetAddr function which can be used to get the net addr from an existing socket (#118)

Author: jjeffcaii

core/transport/tcp_conn.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
+
 	"github.com/rsocket/rsocket-go/core"
 	"github.com/rsocket/rsocket-go/core/framing"
 	"github.com/rsocket/rsocket-go/internal/u24"
@@ -21,6 +22,11 @@ type TCPConn struct {
 	counter *core.TrafficCounter
 }
 
+func (p TCPConn) Addr() string {
+	addr := p.conn.RemoteAddr()
+	return addr.String()
+}
+
 // SetCounter bind a counter which can count r/w bytes.
 func (p *TCPConn) SetCounter(c *core.TrafficCounter) {
 	p.counter = c
@@ -96,7 +102,7 @@ func (p *TCPConn) Write(frame core.WriteableFrame) (err error) {
 	return
 }
 
-// Close close current connection.
+// Close closes current connection.
 func (p *TCPConn) Close() error {
 	return p.conn.Close()
 }

core/transport/tcp_transport.go

@@ -145,12 +145,12 @@ func NewTCPServerTransportWithAddr(network, addr string, tlsConfig *tls.Config)
 	return NewTCPServerTransport(f)
 }
 
-// NewTCPClientTransport creates a new transport.
+// NewTCPClientTransport creates new transport.
 func NewTCPClientTransport(c net.Conn) *Transport {
 	return NewTransport(NewTCPConn(c))
 }
 
-// NewTCPClientTransportWithAddr creates a new transport.
+// NewTCPClientTransportWithAddr creates new transport.
 func NewTCPClientTransportWithAddr(ctx context.Context, network, addr string, tlsConfig *tls.Config) (tp *Transport, err error) {
 	var conn net.Conn
 	var dial net.Dialer

core/transport/transport.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
+
 	"github.com/rsocket/rsocket-go/core"
 	"github.com/rsocket/rsocket-go/core/framing"
 	"github.com/rsocket/rsocket-go/internal/buffer"
@@ -71,7 +72,7 @@ type Transport struct {
 	handlers    [handlerLen]FrameHandler
 }
 
-// NewTransport creates a new transport.
+// NewTransport creates new transport.
 func NewTransport(c Conn) *Transport {
 	return &Transport{
 		conn:        c,
@@ -84,10 +85,18 @@ func IsNoHandlerError(err error) bool {
 	return err == errNoHandler
 }
 
+func (p *Transport) Addr() (string, bool) {
+	ac, ok := p.conn.(AddrConn)
+	if ok {
+		return ac.Addr(), true
+	}
+	return "", false
+}
+
 // Handle register event handlers
 func (p *Transport) Handle(event EventType, handler FrameHandler) {
 	p.mu.Lock()
-	p.handlers[int(event)] = handler
+	p.handlers[event] = handler
 	p.mu.Unlock()
 }
 

core/transport/types.go

@@ -31,6 +31,11 @@ type Conn interface {
 	Read() (core.BufferedFrame, error)
 	// Write writes a frame to Conn.
 	Write(core.WriteableFrame) error
-	// Flush.
+	// Flush flushes the data.
 	Flush() error
 }
+
+type AddrConn interface {
+	Conn
+	Addr() string
+}

core/transport/websocket_conn.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/gorilla/websocket"
 	"github.com/pkg/errors"
+
 	"github.com/rsocket/rsocket-go/core"
 	"github.com/rsocket/rsocket-go/core/framing"
 	"github.com/rsocket/rsocket-go/logger"
@@ -35,6 +36,14 @@ type WebsocketConn struct {
 	counter *core.TrafficCounter
 }
 
+// Addr returns the address info.
+func (p *WebsocketConn) Addr() string {
+	if c, ok := p.c.(*websocket.Conn); ok {
+		return c.RemoteAddr().String()
+	}
+	return ""
+}
+
 // SetCounter bind a counter which can count r/w bytes.
 func (p *WebsocketConn) SetCounter(c *core.TrafficCounter) {
 	p.counter = c

fuzz.go

@@ -1,3 +1,4 @@
+//go:build gofuzz
 // +build gofuzz
 
 //go:generate GO111MODULE=off go-fuzz-build github.com/rsocket/rsocket-go/

internal/socket/base_socket.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
+
 	"github.com/rsocket/rsocket-go/logger"
 	"github.com/rsocket/rsocket-go/payload"
 	"github.com/rsocket/rsocket-go/rx/flux"
@@ -17,6 +18,18 @@ type BaseSocket struct {
 	closers  []func(error)
 	once     sync.Once
 	reqLease *leaser
+	addr     string
+}
+
+func (p *BaseSocket) SetAddr(addr string) {
+	p.addr = addr
+}
+
+func (p *BaseSocket) Addr() (string, bool) {
+	if tp := p.socket.currentTransport(); tp != nil {
+		return tp.Addr()
+	}
+	return p.addr, len(p.addr) > 0
 }
 
 // FireAndForget sends FireAndForget request.

internal/socket/types.go

@@ -54,4 +54,6 @@ type ServerSocket interface {
 	Start(ctx context.Context) error
 	// Token returns token of socket.
 	Token() (token []byte, ok bool)
+	// SetAddr sets address info.
+	SetAddr(addr string)
 }

rsocket.go

@@ -65,6 +65,13 @@ type (
 		RSocket
 	}
 
+	// addressedRSocket is RSocket which contains address info.
+	addressedRSocket interface {
+		RSocket
+		// Addr returns the address info.
+		Addr() (string, bool)
+	}
+
 	// OptAbstractSocket is option for abstract socket.
 	OptAbstractSocket func(*socket.AbstractRSocket)
 )
@@ -113,3 +120,12 @@ func RequestChannel(fn func(requests flux.Flux) (responses flux.Flux)) OptAbstra
 		opts.RC = fn
 	}
 }
+
+// GetAddr returns the address info of given RSocket.
+// Normally, the format is "IP:PORT".
+func GetAddr(rs RSocket) (string, bool) {
+	if ars, ok := rs.(addressedRSocket); ok {
+		return ars.Addr()
+	}
+	return "", false
+}

rsocket_test.go

@@ -14,6 +14,9 @@ import (
 	"github.com/jjeffcaii/reactor-go"
 	"github.com/jjeffcaii/reactor-go/scheduler"
 	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
 	. "github.com/rsocket/rsocket-go"
 	"github.com/rsocket/rsocket-go/core/transport"
 	"github.com/rsocket/rsocket-go/extension"
@@ -23,8 +26,6 @@ import (
 	"github.com/rsocket/rsocket-go/rx"
 	"github.com/rsocket/rsocket-go/rx/flux"
 	"github.com/rsocket/rsocket-go/rx/mono"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
 const (
@@ -313,7 +314,6 @@ func TestSuite(t *testing.T) {
 	for i := 0; i < len(m); i++ {
 		testAll(t, m[i], c[i], s[i])
 	}
-
 }
 
 type ctxKey string