diff --git a/server/connection.go b/server/connection.go index 64947ba5..9f62c9ec 100644 --- a/server/connection.go +++ b/server/connection.go @@ -2,6 +2,7 @@ package server import ( "context" + "net" "github.com/gorilla/websocket" "google.golang.org/protobuf/proto" @@ -16,6 +17,10 @@ type connection struct { var _ types.Connection = (*connection)(nil) +func (c connection) RemoteAddr() net.Addr { + return c.wsConn.RemoteAddr() +} + func (c connection) Send(ctx context.Context, message *protobufs.ServerToAgent) error { bytes, err := proto.Marshal(message) if err != nil { diff --git a/server/serverimpl.go b/server/serverimpl.go index dc34bcac..c2c6eb67 100644 --- a/server/serverimpl.go +++ b/server/serverimpl.go @@ -151,7 +151,7 @@ func (s *server) httpHandler(w http.ResponseWriter, req *http.Request) { } // Return from this func to reduce memory usage. - // Handle the connection on a separate gorountine. + // Handle the connection on a separate goroutine. go s.handleWSConnection(conn) } diff --git a/server/serverimpl_test.go b/server/serverimpl_test.go index 6fabec35..b0da2487 100644 --- a/server/serverimpl_test.go +++ b/server/serverimpl_test.go @@ -89,8 +89,8 @@ func TestServerStartAcceptConnection(t *testing.T) { return types.ConnectionResponse{Accept: true} }, OnConnectedFunc: func(conn types.Connection) { - atomic.StoreInt32(&connectedCalled, 1) srvConn = conn + atomic.StoreInt32(&connectedCalled, 1) }, OnConnectionCloseFunc: func(conn types.Connection) { atomic.StoreInt32(&connectionCloseCalled, 1) @@ -114,6 +114,9 @@ func TestServerStartAcceptConnection(t *testing.T) { eventually(t, func() bool { return atomic.LoadInt32(&connectedCalled) == 1 }) assert.True(t, atomic.LoadInt32(&connectionCloseCalled) == 0) + // Verify that the RemoteAddr is correct. + require.Equal(t, conn.LocalAddr().String(), srvConn.RemoteAddr().String()) + // Close the connection from client side. conn.Close() diff --git a/server/types/connection.go b/server/types/connection.go index fe706be3..26da474d 100644 --- a/server/types/connection.go +++ b/server/types/connection.go @@ -2,6 +2,7 @@ package types import ( "context" + "net" "github.com/open-telemetry/opamp-go/protobufs" ) @@ -9,6 +10,9 @@ import ( // Connection represents one OpAMP WebSocket connections. // The implementation MUST be a comparable type so that it can be used as a map key. type Connection interface { + // RemoteAddr returns the remote network address of the connection. + RemoteAddr() net.Addr + // Send a message. Should not be called concurrently for the same Connection instance. // Blocks until the message is sent. // Should return as soon as possible if the ctx is cancelled.