Repository: open-telemetry/opamp-go
Pull request #153 (2022-11-29 23:22:19)

Change Attach signature to return ConnContext func
Resolves https://github.com/open-telemetry/opamp-go/issues/152

Make sure the library-provided ConnContext function is used by the caller when they setup the HTTP Server to Attach() to, so that the context correctly stores a references to the connection. This is an internal requirement for plain http server implementation.


linux-foundation-easycla[bot]:  <a href="https://api.easycla.lfx.linuxfoundation.org/v2/repository-provider/github/sign/9522512/424240089/153/#/?version=2"><img src="https://s3.amazonaws.com/cla-project-logo-prod/cla-not-signed.svg" alt="CLA Not Signed" align="left" height="28" width="328"></a><br/><br /><ul><li><a href='https://api.easycla.lfx.linuxfoundation.org/v2/repository-provider/github/sign/9522512/424240089/153/#/?version=2' target='_blank'>:x:</a> - login: @tshinde-splunk . The commit (b4f6e3f8a4b1c54237e4e0425dc4b659bf92a395) is not authorized under a signed CLA. <a href='https://api.easycla.lfx.linuxfoundation.org/v2/repository-provider/github/sign/9522512/424240089/153/#/?version=2' target='_blank'>Please click here to be authorized</a>. For further assistance with EasyCLA, <a href='https://jira.linuxfoundation.org/servicedesk/customer/portal/4' target='_blank'>please submit a support request ticket</a>.</li></ul>

Patch:
diff --git a/client/httpclient.go b/client/httpclient.go
index fea31a9..0fc34aa 100644
--- a/client/httpclient.go
+++ b/client/httpclient.go
@@ -45,6 +45,9 @@ func (c *httpClient) Start(ctx context.Context, settings types.StartSettings) er
 	// Prepare Server connection settings.
 	c.sender.SetRequestHeader(settings.Header)
 
+	// Add TLS configuration into httpClient
+	c.sender.AddTLSConfig(settings.TLSConfig)
+
 	if settings.EnableCompression {
 		c.sender.EnableCompression()
 	}
diff --git a/client/internal/httpsender.go b/client/internal/httpsender.go
index b2a9938..42244b4 100644
--- a/client/internal/httpsender.go
+++ b/client/internal/httpsender.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"compress/gzip"
 	"context"
+	"crypto/tls"
 	"errors"
 	"fmt"
 	"io"
@@ -263,3 +264,11 @@ func (h *HTTPSender) EnableCompression() {
 	h.compressionEnabled = true
 	h.requestHeader.Set(headerContentEncoding, encodingTypeGZip)
 }
+
+func (h *HTTPSender) AddTLSConfig(config *tls.Config) {
+	if config != nil {
+		h.client.Transport = &http.Transport{
+			TLSClientConfig: config,
+		}
+	}
+}
diff --git a/client/internal/mockserver.go b/client/internal/mockserver.go
index 11a7013..6a85ec6 100644
--- a/client/internal/mockserver.go
+++ b/client/internal/mockserver.go
@@ -37,7 +37,7 @@ type MockServer struct {
 const headerContentType = "Content-Type"
 const contentTypeProtobuf = "application/x-protobuf"
 
-func StartMockServer(t *testing.T) *MockServer {
+func newMockServer(t *testing.T) (*MockServer, *http.ServeMux) {
 	srv := &MockServer{
 		t:                t,
 		expectedHandlers: make(chan receivedMessageHandler, 0),
@@ -65,6 +65,12 @@ func StartMockServer(t *testing.T) *MockServer {
 		},
 	)
 
+	return srv, m
+}
+
+func StartMockServer(t *testing.T) *MockServer {
+	srv, m := newMockServer(t)
+
 	srv.srv = httptest.NewServer(m)
 
 	u, err := url.Parse(srv.srv.URL)
@@ -78,6 +84,26 @@ func StartMockServer(t *testing.T) *MockServer {
 	return srv
 }
 
+func StartTLSMockServer(t *testing.T) *MockServer {
+	srv, m := newMockServer(t)
+
+	srv.srv = httptest.NewTLSServer(m)
+
+	u, err := url.Parse(srv.srv.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	srv.Endpoint = u.Host
+
+	testhelpers.WaitForEndpoint(srv.Endpoint)
+
+	return srv
+}
+
+func (m *MockServer) GetHTTPTestServer() *httptest.Server {
+	return m.srv
+}
+
 // EnableExpectMode enables the expect mode that allows using Expect() method
 // to describe what message is expected to be received.
 func (m *MockServer) EnableExpectMode() {
diff --git a/server/server.go b/server/server.go
index d28f9b6..4403304 100644
--- a/server/server.go
+++ b/server/server.go
@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"crypto/tls"
+	"net"
 	"net/http"
 
 	"github.com/open-telemetry/opamp-go/server/types"
@@ -34,18 +35,21 @@ type StartSettings struct {
 
 type HTTPHandlerFunc func(http.ResponseWriter, *http.Request)
 
+type ConnContext func(ctx context.Context, c net.Conn) context.Context
+
 type OpAMPServer interface {
 	// Attach prepares the OpAMP Server to begin handling requests from an existing
-	// http.Server. The returned HTTPHandlerFunc should be added as a handler to the
-	// desired http.Server by the caller and the http.Server should be started by
-	// the caller after that.
+	// http.Server. The returned HTTPHandlerFunc and ConnContext should be added as a
+	// handler and ConnContext respectively to the desired http.Server by the caller
+	// and the http.Server should be started by the caller after that. The ConnContext
+	// is only used for plain http connections.
 	// For example:
-	//   handler, _ := Server.Attach()
+	//   handler, connContext, _ := Server.Attach()
 	//   mux := http.NewServeMux()
 	//   mux.HandleFunc("/opamp", handler)
-	//   httpSrv := &http.Server{Handler:mux,Addr:"127.0.0.1:4320"}
+	//   httpSrv := &http.Server{Handler:mux,Addr:"127.0.0.1:4320", ConnContext: connContext}
 	//   httpSrv.ListenAndServe()
-	Attach(settings Settings) (HTTPHandlerFunc, error)
+	Attach(settings Settings) (HTTPHandlerFunc, ConnContext, error)
 
 	// Start an OpAMP Server and begin accepting connections. Starts its own http.Server
 	// using provided settings. This should block until the http.Server is ready to
diff --git a/server/serverimpl.go b/server/serverimpl.go
index 8826406..9758acd 100644
--- a/server/serverimpl.go
+++ b/server/serverimpl.go
@@ -50,12 +50,12 @@ func New(logger types.Logger) *server {
 	return &server{logger: logger}
 }
 
-func (s *server) Attach(settings Settings) (HTTPHandlerFunc, error) {
+func (s *server) Attach(settings Settings) (HTTPHandlerFunc, ConnContext, error) {
 	s.settings = settings
 	s.wsUpgrader = websocket.Upgrader{
 		EnableCompression: settings.EnableCompression,
 	}
-	return s.httpHandler, nil
+	return s.httpHandler, contextWithConn, nil
 }
 
 func (s *server) Start(settings StartSettings) error {
@@ -63,7 +63,7 @@ func (s *server) Start(settings StartSettings) error {
 		return errAlreadyStarted
 	}
 
-	_, err := s.Attach(settings.Settings)
+	_, _, err := s.Attach(settings.Settings)
 	if err != nil {
 		return err
 	}


Test patch:
diff --git a/client/clientimpl_test.go b/client/clientimpl_test.go
index 2b23751..576244f 100644
--- a/client/clientimpl_test.go
+++ b/client/clientimpl_test.go
@@ -2,6 +2,8 @@ package client
 
 import (
 	"context"
+	"crypto/tls"
+	"crypto/x509"
 	"math/rand"
 	"net/http"
 	"net/http/httptest"
@@ -83,9 +85,17 @@ func prepareSettings(t *testing.T, settings *types.StartSettings, c OpAMPClient)
 	require.NoError(t, err)
 	switch c.(type) {
 	case *httpClient:
-		u.Scheme = "http"
+		if settings.TLSConfig != nil {
+			u.Scheme = "https"
+		} else {
+			u.Scheme = "http"
+		}
 	case *wsClient:
-		u.Scheme = "ws"
+		if settings.TLSConfig != nil {
+			u.Scheme = "wss"
+		} else {
+			u.Scheme = "ws"
+		}
 	}
 	settings.OpAMPServerURL = u.String()
 }
@@ -305,6 +315,50 @@ func TestConnectWithHeader(t *testing.T) {
 	})
 }
 
+func TestConnectWithTLS(t *testing.T) {
+	testClients(t, func(t *testing.T, client OpAMPClient) {
+		// Start a server.
+		srv := internal.StartTLSMockServer(t)
+		var conn atomic.Value
+		srv.OnConnect = func(r *http.Request) {
+			conn.Store(true)
+		}
+
+		certs := rootCAs(t, srv.GetHTTPTestServer())
+
+		// Start a client.
+		settings := types.StartSettings{
+			OpAMPServerURL: "wss://" + srv.Endpoint,
+			TLSConfig: &tls.Config{
+				RootCAs: certs,
+			},
+		}
+
+		startClient(t, settings, client)
+
+		// Wait for connection to be established.
+		eventually(t, func() bool { return conn.Load() != nil })
+
+		// Shutdown the Server and the client.
+		srv.Close()
+		_ = client.Stop(context.Background())
+	})
+}
+
+func rootCAs(t *testing.T, s *httptest.Server) *x509.CertPool {
+	certs := x509.NewCertPool()
+	for _, c := range s.TLS.Certificates {
+		roots, err := x509.ParseCertificates(c.Certificate[len(c.Certificate)-1])
+		if err != nil {
+			t.Fatalf("error parsing server's root cert: %v", err)
+		}
+		for _, root := range roots {
+			certs.AddCert(root)
+		}
+	}
+	return certs
+}
+
 func createRemoteConfig() *protobufs.AgentRemoteConfig {
 	return &protobufs.AgentRemoteConfig{
 		Config: &protobufs.AgentConfigMap{
diff --git a/client/internal/httpsender_test.go b/client/internal/httpsender_test.go
index 31f1aab..58a5268 100644
--- a/client/internal/httpsender_test.go
+++ b/client/internal/httpsender_test.go
@@ -2,6 +2,7 @@ package internal
 
 import (
 	"context"
+	"crypto/tls"
 	"net/http"
 	"sync/atomic"
 	"testing"
@@ -55,3 +56,45 @@ func TestHTTPSenderRetryForStatusTooManyRequests(t *testing.T) {
 	cancel()
 	srv.Close()
 }
+
+func TestAddTLSConfig(t *testing.T) {
+	sender := NewHTTPSender(&sharedinternal.NopLogger{})
+
+	certificate, err := GenerateCertificate()
+	assert.NoError(t, err)
+
+	tlsConfig := &tls.Config{
+		Certificates: []tls.Certificate{certificate},
+	}
+
+	sender.AddTLSConfig(tlsConfig)
+	assert.Equal(t, sender.client.Transport, &http.Transport{TLSClientConfig: tlsConfig})
+}
+
+func GenerateCertificate() (tls.Certificate, error) {
+
+	certPem := []byte(`-----BEGIN CERTIFICATE-----
+MIIBhTCCASugAwIBAgIQIRi6zePL6mKjOipn+dNuaTAKBggqhkjOPQQDAjASMRAw
+DgYDVQQKEwdBY21lIENvMB4XDTE3MTAyMDE5NDMwNloXDTE4MTAyMDE5NDMwNlow
+EjEQMA4GA1UEChMHQWNtZSBDbzBZMBMGByqGSM49AgEGCCqGSM49AwEHA0IABD0d
+7VNhbWvZLWPuj/RtHFjvtJBEwOkhbN/BnnE8rnZR8+sbwnc/KhCk3FhnpHZnQz7B
+5aETbbIgmuvewdjvSBSjYzBhMA4GA1UdDwEB/wQEAwICpDATBgNVHSUEDDAKBggr
+BgEFBQcDATAPBgNVHRMBAf8EBTADAQH/MCkGA1UdEQQiMCCCDmxvY2FsaG9zdDo1
+NDUzgg4xMjcuMC4wLjE6NTQ1MzAKBggqhkjOPQQDAgNIADBFAiEA2zpJEPQyz6/l
+Wf86aX6PepsntZv2GYlA5UpabfT2EZICICpJ5h/iI+i341gBmLiAFQOyTDT+/wQc
+6MF9+Yw1Yy0t
+-----END CERTIFICATE-----`)
+
+	keyPem := []byte(`-----BEGIN EC PRIVATE KEY-----
+MHcCAQEEIIrYSSNQFaA2Hwf1duRSxKtLYX5CB04fSeQ6tF1aY/PuoAoGCCqGSM49
+AwEHoUQDQgAEPR3tU2Fta9ktY+6P9G0cWO+0kETA6SFs38GecTyudlHz6xvCdz8q
+EKTcWGekdmdDPsHloRNtsiCa697B2O9IFA==
+-----END EC PRIVATE KEY-----`)
+
+	cert, err := tls.X509KeyPair(certPem, keyPem)
+	if err != nil {
+		return tls.Certificate{}, err
+	}
+
+	return cert, nil
+}
diff --git a/server/serverimpl_test.go b/server/serverimpl_test.go
index 3b0749f..129518e 100644
--- a/server/serverimpl_test.go
+++ b/server/serverimpl_test.go
@@ -422,7 +422,7 @@ func TestServerAttachAcceptConnection(t *testing.T) {
 	settings := Settings{Callbacks: callbacks}
 	srv := New(&sharedinternal.NopLogger{})
 	require.NotNil(t, srv)
-	handlerFunc, err := srv.Attach(settings)
+	handlerFunc, _, err := srv.Attach(settings)
 	require.NoError(t, err)
 
 	// Create an HTTP Server and make it handle OpAMP connections.
@@ -447,6 +447,88 @@ func TestServerAttachAcceptConnection(t *testing.T) {
 	eventually(t, func() bool { return atomic.LoadInt32(&connectionCloseCalled) == 1 })
 }
 
+func TestServerAttachSendMessagePlainHTTP(t *testing.T) {
+	connectedCalled := int32(0)
+	connectionCloseCalled := int32(0)
+	var rcvMsg atomic.Value
+
+	var srvConn types.Connection
+	callbacks := CallbacksStruct{
+		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
+			return types.ConnectionResponse{Accept: true}
+		},
+		OnConnectedFunc: func(conn types.Connection) {
+			atomic.StoreInt32(&connectedCalled, 1)
+			srvConn = conn
+		},
+		OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
+			// Remember received message.
+			rcvMsg.Store(message)
+
+			// Send a response.
+			response := protobufs.ServerToAgent{
+				InstanceUid:  message.InstanceUid,
+				Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
+			}
+			return &response
+		},
+		OnConnectionCloseFunc: func(conn types.Connection) {
+			atomic.StoreInt32(&connectionCloseCalled, 1)
+			assert.EqualValues(t, srvConn, conn)
+		},
+	}
+
+	// Prepare to attach OpAMP Server to an HTTP Server created separately.
+	settings := Settings{Callbacks: callbacks}
+	srv := New(&sharedinternal.NopLogger{})
+	require.NotNil(t, srv)
+	handlerFunc, ContextWithConn, err := srv.Attach(settings)
+	require.NoError(t, err)
+
+	// Create an HTTP Server and make it handle OpAMP connections.
+	mux := http.NewServeMux()
+	path := "/opamppath"
+	mux.HandleFunc(path, handlerFunc)
+	hs := httptest.NewUnstartedServer(mux)
+	hs.Config.ConnContext = ContextWithConn
+	hs.Start()
+	defer hs.Close()
+
+	// Send a message to the Server.
+	sendMsg := protobufs.AgentToServer{
+		InstanceUid: "12345678",
+	}
+	b, err := proto.Marshal(&sendMsg)
+	require.NoError(t, err)
+	resp, err := http.Post("http://"+hs.Listener.Addr().String()+path, contentTypeProtobuf, bytes.NewReader(b))
+	require.NoError(t, err)
+
+	// Wait until Server receives the message.
+	eventually(t, func() bool { return rcvMsg.Load() != nil })
+	assert.True(t, atomic.LoadInt32(&connectedCalled) == 1)
+
+	// Verify the received message is what was sent.
+	assert.True(t, proto.Equal(rcvMsg.Load().(proto.Message), &sendMsg))
+
+	// Read Server's response.
+	b, err = io.ReadAll(resp.Body)
+	require.NoError(t, err)
+
+	assert.EqualValues(t, http.StatusOK, resp.StatusCode)
+	assert.EqualValues(t, contentTypeProtobuf, resp.Header.Get(headerContentType))
+
+	// Decode the response.
+	var response protobufs.ServerToAgent
+	err = proto.Unmarshal(b, &response)
+	require.NoError(t, err)
+
+	// Verify the response.
+	assert.EqualValues(t, sendMsg.InstanceUid, response.InstanceUid)
+	assert.EqualValues(t, protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus, response.Capabilities)
+
+	eventually(t, func() bool { return atomic.LoadInt32(&connectionCloseCalled) == 1 })
+}
+
 func TestServerHonoursClientRequestContentEncoding(t *testing.T) {
 
 	hc := http.Client{}