Repository: open-telemetry/opamp-go (go)

Issue:
Move certs helper code from /internal/examples to /internal or /internal/testhelper?
In an effort to improve the test coverage for [serverimpl](https://app.codecov.io/gh/open-telemetry/opamp-go/blob/main/server%2Fserverimpl.go), I started to add some tests with TLS enabled. I noticed we already have some useful code for generating and reading certs in `internal/examples/...`, and it seems like a good idea to use for tests. I was thinking of moving them to `/internal` or maybe `internal/testhelper` and use in both serverimpl tests and examples.

@tigrannajaryan do you see any issues with this? If not, I can send a PR for this along with the new serverimpl TLS tests.

Pull request:
Move certs & helper code from /internal/examples to /internal
Fixes #199

I verified that examples run with the changes.

Patch:
diff --git a/go.mod b/go.mod
index 134b96a..93ae25c 100644
--- a/go.mod
+++ b/go.mod
@@ -1,6 +1,6 @@
 module github.com/open-telemetry/opamp-go
 
-go 1.17
+go 1.20
 
 require (
 	github.com/cenkalti/backoff/v4 v4.1.2
diff --git a/internal/certs.go b/internal/certs.go
new file mode 100644
index 0000000..7b86e07
--- /dev/null
+++ b/internal/certs.go
@@ -0,0 +1,164 @@
+package internal
+
+// This file contains helper functions to read and
+// create TLS configs/certificates. Currently used in
+// the example client and server and
+// in the tests. Not intended for any other use.
+
+import (
+	"bytes"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/tls"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"encoding/pem"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"math/big"
+	"net"
+	"os"
+	"time"
+
+	"github.com/open-telemetry/opamp-go/protobufs"
+)
+
+func CreateClientTLSConfig(clientCert *tls.Certificate, caCertPath string) (*tls.Config, error) {
+	// Read the CA's public key. This is the CA that signs the server's certificate.
+	caCertBytes, err := os.ReadFile(caCertPath)
+	if err != nil {
+		return nil, err
+	}
+
+	// Create a certificate pool and make our CA trusted.
+	caCertPool := x509.NewCertPool()
+	if ok := caCertPool.AppendCertsFromPEM(caCertBytes); !ok {
+		return nil, errors.New("cannot append ca.cert.pem")
+	}
+
+	cfg := &tls.Config{
+		RootCAs: caCertPool,
+	}
+	if clientCert != nil {
+		// If there is a client-side certificate use it for connection too.
+		cfg.Certificates = []tls.Certificate{*clientCert}
+	}
+	return cfg, nil
+}
+
+func CreateServerTLSConfig(caCertPath, serverCertPath, serverKeyPath string) (*tls.Config, error) {
+	// Read the CA's public key. This is the CA that signs the server's certificate.
+	caCertBytes, err := os.ReadFile(caCertPath)
+	if err != nil {
+		return nil, err
+	}
+
+	// Create a certificate pool and make our CA trusted.
+	caCertPool := x509.NewCertPool()
+	if ok := caCertPool.AppendCertsFromPEM(caCertBytes); !ok {
+		return nil, errors.New("cannot append ca.cert.pem")
+	}
+
+	// Load server's certificate.
+	cert, err := tls.LoadX509KeyPair(
+		serverCertPath,
+		serverKeyPath,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("tls.LoadX509KeyPair failed: %v", err)
+	}
+	tlsConfig := &tls.Config{
+		Certificates: []tls.Certificate{cert},
+		// TODO: verify client cert manually, and allow TOFU option. See manual
+		// verification example: https://dev.to/living_syn/validating-client-certificate-sans-in-go-i5p
+		// Instead, we use VerifyClientCertIfGiven which will automatically verify the provided certificate
+		// is signed by our CA (so TOFU with self-generated client certificate will not work).
+		ClientAuth: tls.VerifyClientCertIfGiven,
+		// Allow insecure connections for demo purposes.
+		InsecureSkipVerify: true,
+		ClientCAs:          caCertPool,
+	}
+	tlsConfig.BuildNameToCertificate()
+	return tlsConfig, nil
+}
+
+func CreateTLSCert(caCertPath, caKeyPath string) (*protobufs.TLSCertificate, error) {
+
+	// Load CA Cert.
+	caCertBytes, err := ioutil.ReadFile(caCertPath)
+	if err != nil {
+		return nil, fmt.Errorf("cannot read CA cert: %v", err)
+	}
+
+	caKeyBytes, err := ioutil.ReadFile(caKeyPath)
+	if err != nil {
+		return nil, fmt.Errorf("cannot read CA key: %v", err)
+	}
+
+	caCertPB, _ := pem.Decode(caCertBytes)
+	caKeyPB, _ := pem.Decode(caKeyBytes)
+	caCert, err := x509.ParseCertificate(caCertPB.Bytes)
+	if err != nil {
+		return nil, fmt.Errorf("cannot parse CA cert: %v", err)
+	}
+
+	caPrivKey, err := x509.ParsePKCS1PrivateKey(caKeyPB.Bytes)
+	if err != nil {
+		return nil, fmt.Errorf("cannot parse CA key: %v", err)
+	}
+
+	// Generate a private key for new client cert.
+	certPrivKey, err := rsa.GenerateKey(rand.Reader, 4096)
+	if err != nil {
+		err := fmt.Errorf("cannot generate private key: %v", err)
+		return nil, err
+	}
+
+	// Prepare certificate template.
+	template := &x509.Certificate{
+		SerialNumber: big.NewInt(1),
+		Subject: pkix.Name{
+			CommonName:    "OpAMP Example Client",
+			Organization:  []string{"OpAMP Example"},
+			Country:       []string{"CA"},
+			Province:      []string{"ON"},
+			Locality:      []string{"City"},
+			StreetAddress: []string{""},
+			PostalCode:    []string{""},
+		},
+		IPAddresses: []net.IP{net.IPv4(127, 0, 0, 1)},
+		NotBefore:   time.Now(),
+		NotAfter:    time.Now().Add(time.Hour * 1000),
+		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
+		KeyUsage:    x509.KeyUsageDigitalSignature,
+	}
+
+	// Create the client cert. Sign it using CA cert.
+	certBytes, err := x509.CreateCertificate(rand.Reader, template, caCert, &certPrivKey.PublicKey, caPrivKey)
+	if err != nil {
+		err := fmt.Errorf("cannot create certificate: %v", err)
+		return nil, err
+	}
+
+	publicKeyPEM := new(bytes.Buffer)
+	pem.Encode(publicKeyPEM, &pem.Block{
+		Type:  "CERTIFICATE",
+		Bytes: certBytes,
+	})
+
+	privateKeyPEM := new(bytes.Buffer)
+	pem.Encode(privateKeyPEM, &pem.Block{
+		Type:  "RSA PRIVATE KEY",
+		Bytes: x509.MarshalPKCS1PrivateKey(certPrivKey),
+	})
+
+	// We have a client certificate with a public and private key.
+	certificate := &protobufs.TLSCertificate{
+		PublicKey:   publicKeyPEM.Bytes(),
+		PrivateKey:  privateKeyPEM.Bytes(),
+		CaPublicKey: caCertBytes,
+	}
+
+	return certificate, nil
+}
diff --git a/internal/examples/certs/certs/ca.cert.pem b/internal/certs/certs/ca.cert.pem
similarity index 100%
rename from internal/examples/certs/certs/ca.cert.pem
rename to internal/certs/certs/ca.cert.pem
diff --git a/internal/examples/certs/clear.sh b/internal/certs/clear.sh
similarity index 100%
rename from internal/examples/certs/clear.sh
rename to internal/certs/clear.sh
diff --git a/internal/examples/certs/client.conf b/internal/certs/client.conf
similarity index 100%
rename from internal/examples/certs/client.conf
rename to internal/certs/client.conf
diff --git a/internal/examples/certs/generate.sh b/internal/certs/generate.sh
similarity index 100%
rename from internal/examples/certs/generate.sh
rename to internal/certs/generate.sh
diff --git a/internal/examples/certs/index.txt b/internal/certs/index.txt
similarity index 100%
rename from internal/examples/certs/index.txt
rename to internal/certs/index.txt
diff --git a/internal/examples/certs/index.txt.attr b/internal/certs/index.txt.attr
similarity index 100%
rename from internal/examples/certs/index.txt.attr
rename to internal/certs/index.txt.attr
diff --git a/internal/examples/certs/index.txt.old b/internal/certs/index.txt.old
similarity index 100%
rename from internal/examples/certs/index.txt.old
rename to internal/certs/index.txt.old
diff --git a/internal/examples/certs/openssl.conf b/internal/certs/openssl.conf
similarity index 100%
rename from internal/examples/certs/openssl.conf
rename to internal/certs/openssl.conf
diff --git a/internal/examples/certs/private/ca.key.pem b/internal/certs/private/ca.key.pem
similarity index 100%
rename from internal/examples/certs/private/ca.key.pem
rename to internal/certs/private/ca.key.pem
diff --git a/internal/examples/certs/serial b/internal/certs/serial
similarity index 100%
rename from internal/examples/certs/serial
rename to internal/certs/serial
diff --git a/internal/examples/certs/serial.old b/internal/certs/serial.old
similarity index 100%
rename from internal/examples/certs/serial.old
rename to internal/certs/serial.old
diff --git a/internal/examples/certs/server.conf b/internal/certs/server.conf
similarity index 100%
rename from internal/examples/certs/server.conf
rename to internal/certs/server.conf
diff --git a/internal/examples/certs/server_certs/server.cert.pem b/internal/certs/server_certs/server.cert.pem
similarity index 100%
rename from internal/examples/certs/server_certs/server.cert.pem
rename to internal/certs/server_certs/server.cert.pem
diff --git a/internal/examples/certs/server_certs/server.csr b/internal/certs/server_certs/server.csr
similarity index 100%
rename from internal/examples/certs/server_certs/server.csr
rename to internal/certs/server_certs/server.csr
diff --git a/internal/examples/certs/server_certs/server.key.pem b/internal/certs/server_certs/server.key.pem
similarity index 100%
rename from internal/examples/certs/server_certs/server.key.pem
rename to internal/certs/server_certs/server.key.pem
diff --git a/internal/examples/certs/server_ext.conf b/internal/certs/server_ext.conf
similarity index 100%
rename from internal/examples/certs/server_ext.conf
rename to internal/certs/server_ext.conf
diff --git a/internal/examples/agent/agent/agent.go b/internal/examples/agent/agent/agent.go
index c82d96c..5e0a496 100644
--- a/internal/examples/agent/agent/agent.go
+++ b/internal/examples/agent/agent/agent.go
@@ -6,7 +6,6 @@ import (
 	"crypto/x509"
 	"encoding/pem"
 	"fmt"
-	"log"
 	"math/rand"
 	"os"
 	"runtime"
@@ -20,6 +19,7 @@ import (
 
 	"github.com/open-telemetry/opamp-go/client"
 	"github.com/open-telemetry/opamp-go/client/types"
+	"github.com/open-telemetry/opamp-go/internal"
 	"github.com/open-telemetry/opamp-go/protobufs"
 )
 
@@ -89,9 +89,17 @@ func NewAgent(logger types.Logger, agentType string, agentVersion string) *Agent
 func (agent *Agent) connect() error {
 	agent.opampClient = client.NewWebSocket(agent.logger)
 
+	tlsConfig, err := internal.CreateClientTLSConfig(
+		agent.opampClientCert,
+		"../../certs/certs/ca.cert.pem",
+	)
+	if err != nil {
+		return err
+	}
+
 	settings := types.StartSettings{
 		OpAMPServerURL: "wss://127.0.0.1:4320/v1/opamp",
-		TLSConfig:      createClientTLSConfig(agent.opampClientCert),
+		TLSConfig:      tlsConfig,
 		InstanceUid:    agent.instanceId.String(),
 		Callbacks: types.CallbacksStruct{
 			OnConnectFunc: func() {
@@ -120,7 +128,7 @@ func (agent *Agent) connect() error {
 			protobufs.AgentCapabilities_AgentCapabilities_AcceptsOpAMPConnectionSettings,
 	}
 
-	err := agent.opampClient.SetAgentDescription(agent.agentDescription)
+	err = agent.opampClient.SetAgentDescription(agent.agentDescription)
 	if err != nil {
 		return err
 	}
@@ -137,29 +145,6 @@ func (agent *Agent) connect() error {
 	return nil
 }
 
-func createClientTLSConfig(clientCert *tls.Certificate) *tls.Config {
-	// Read the CA's public key. This is the CA that signs the server's certificate.
-	caCertBytes, err := os.ReadFile("../certs/certs/ca.cert.pem")
-	if err != nil {
-		log.Fatalln(err)
-	}
-
-	// Create a certificate pool and make our CA trusted.
-	caCertPool := x509.NewCertPool()
-	if ok := caCertPool.AppendCertsFromPEM(caCertBytes); !ok {
-		log.Fatalln("Cannot append ca.cert.pem")
-	}
-
-	cfg := &tls.Config{
-		RootCAs: caCertPool,
-	}
-	if clientCert != nil {
-		// If there is a client-side certificate use it for connection too.
-		cfg.Certificates = []tls.Certificate{*clientCert}
-	}
-	return cfg
-}
-
 func (agent *Agent) disconnect() {
 	agent.logger.Debugf("Disconnecting from server...")
 	agent.opampClient.Stop(context.Background())
diff --git a/internal/examples/go.mod b/internal/examples/go.mod
index f79cc48..d3cfb4f 100644
--- a/internal/examples/go.mod
+++ b/internal/examples/go.mod
@@ -1,12 +1,14 @@
 module github.com/open-telemetry/opamp-go/internal/examples
 
-go 1.17
+go 1.20
 
 require (
+	github.com/cenkalti/backoff/v4 v4.1.2
 	github.com/knadh/koanf v1.3.3
 	github.com/oklog/ulid/v2 v2.0.2
 	github.com/open-telemetry/opamp-go v0.1.0
 	github.com/shirou/gopsutil v3.21.11+incompatible
+	github.com/stretchr/testify v1.7.0
 	go.opentelemetry.io/otel v1.3.0
 	go.opentelemetry.io/otel/exporters/otlp/otlpmetric v0.26.0
 	go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp v0.26.0
@@ -17,7 +19,6 @@ require (
 )
 
 require (
-	github.com/cenkalti/backoff/v4 v4.1.2 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fsnotify/fsnotify v1.4.9 // indirect
 	github.com/go-logr/logr v1.2.1 // indirect
@@ -30,7 +31,6 @@ require (
 	github.com/mitchellh/mapstructure v1.4.1 // indirect
 	github.com/mitchellh/reflectwalk v1.0.2 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/stretchr/testify v1.7.0 // indirect
 	github.com/tklauser/go-sysconf v0.3.9 // indirect
 	github.com/tklauser/numcpus v0.3.0 // indirect
 	github.com/yusufpapurcu/wmi v1.2.2 // indirect
diff --git a/internal/examples/server/data/agent.go b/internal/examples/server/data/agent.go
index e920f03..b562bd3 100644
--- a/internal/examples/server/data/agent.go
+++ b/internal/examples/server/data/agent.go
@@ -25,8 +25,6 @@ type Agent struct {
 
 	// Connection to the Agent.
 	conn types.Connection
-	// Mutex to protect Send() operation.
-	connMutex sync.Mutex
 
 	// mutex for the fields that follow it.
 	mux sync.RWMutex
@@ -421,9 +419,6 @@ func (agent *Agent) calcConnectionSettings(response *protobufs.ServerToAgent) {
 }
 
 func (agent *Agent) SendToAgent(msg *protobufs.ServerToAgent) {
-	agent.connMutex.Lock()
-	defer agent.connMutex.Unlock()
-
 	agent.conn.Send(context.Background(), msg)
 }
 
diff --git a/internal/examples/server/opampsrv/opampsrv.go b/internal/examples/server/opampsrv/opampsrv.go
index 908b10a..4beb781 100644
--- a/internal/examples/server/opampsrv/opampsrv.go
+++ b/internal/examples/server/opampsrv/opampsrv.go
@@ -2,15 +2,10 @@ package opampsrv
 
 import (
 	"context"
-	"crypto/tls"
-	"crypto/x509"
-	"errors"
-	"fmt"
 	"log"
 	"net/http"
-	"os"
-	"path"
 
+	"github.com/open-telemetry/opamp-go/internal"
 	"github.com/open-telemetry/opamp-go/internal/examples/server/data"
 	"github.com/open-telemetry/opamp-go/protobufs"
 	"github.com/open-telemetry/opamp-go/server"
@@ -59,7 +54,11 @@ func (srv *Server) Start() {
 		},
 		ListenEndpoint: "127.0.0.1:4320",
 	}
-	tlsConfig, err := createServerTLSConfig("../certs")
+	tlsConfig, err := internal.CreateServerTLSConfig(
+		"../../certs/certs/ca.cert.pem",
+		"../../certs/server_certs/server.cert.pem",
+		"../../certs/server_certs/server.key.pem",
+	)
 	if err != nil {
 		srv.logger.Debugf("Could not load TLS config, working without TLS: %v", err.Error())
 	}
@@ -68,42 +67,6 @@ func (srv *Server) Start() {
 	srv.opampSrv.Start(settings)
 }
 
-func createServerTLSConfig(certsDir string) (*tls.Config, error) {
-	// Read the CA's public key. This is the CA that signs the server's certificate.
-	caCertBytes, err := os.ReadFile(path.Join(certsDir, "certs/ca.cert.pem"))
-	if err != nil {
-		return nil, err
-	}
-
-	// Create a certificate pool and make our CA trusted.
-	caCertPool := x509.NewCertPool()
-	if ok := caCertPool.AppendCertsFromPEM(caCertBytes); !ok {
-		return nil, errors.New("cannot append ca.cert.pem")
-	}
-
-	// Load server's certificate.
-	cert, err := tls.LoadX509KeyPair(
-		path.Join(certsDir, "server_certs/server.cert.pem"),
-		path.Join(certsDir, "server_certs/server.key.pem"),
-	)
-	if err != nil {
-		return nil, fmt.Errorf("tls.LoadX509KeyPair failed: %v", err)
-	}
-	tlsConfig := &tls.Config{
-		Certificates: []tls.Certificate{cert},
-		// TODO: verify client cert manually, and allow TOFU option. See manual
-		// verification example: https://dev.to/living_syn/validating-client-certificate-sans-in-go-i5p
-		// Instead, we use VerifyClientCertIfGiven which will automatically verify the provided certificate
-		// is signed by our CA (so TOFU with self-generated client certificate will not work).
-		ClientAuth: tls.VerifyClientCertIfGiven,
-		// Allow insecure connections for demo purposes.
-		InsecureSkipVerify: true,
-		ClientCAs:          caCertPool,
-	}
-	tlsConfig.BuildNameToCertificate()
-	return tlsConfig, nil
-}
-
 func (srv *Server) Stop() {
 	srv.opampSrv.Stop(context.Background())
 }
diff --git a/internal/examples/server/uisrv/ui.go b/internal/examples/server/uisrv/ui.go
index ec29628..3e18f48 100644
--- a/internal/examples/server/uisrv/ui.go
+++ b/internal/examples/server/uisrv/ui.go
@@ -1,23 +1,14 @@
 package uisrv
 
 import (
-	"bytes"
 	"context"
-	"crypto/rand"
-	"crypto/rsa"
-	"crypto/x509"
-	"crypto/x509/pkix"
-	"encoding/pem"
-	"fmt"
-	"io/ioutil"
 	"log"
-	"math/big"
-	"net"
 	"net/http"
 	"path"
 	"text/template"
 	"time"
 
+	"github.com/open-telemetry/opamp-go/internal"
 	"github.com/open-telemetry/opamp-go/internal/examples/server/data"
 	"github.com/open-telemetry/opamp-go/protobufs"
 )
@@ -130,7 +121,7 @@ func rotateInstanceClientCert(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Create a new certificate for the agent.
-	certificate, err := createTLSCert()
+	certificate, err := internal.CreateTLSCert("../../certs/certs/ca.cert.pem", "../../certs/certs/ca.key.pem")
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
 		logger.Println(err)
@@ -162,84 +153,3 @@ func rotateInstanceClientCert(w http.ResponseWriter, r *http.Request) {
 
 	http.Redirect(w, r, "/agent?instanceid="+string(instanceId), http.StatusSeeOther)
 }
-
-func createTLSCert() (*protobufs.TLSCertificate, error) {
-	certsDir := "../certs"
-
-	// Load CA Cert.
-	caCertBytes, err := ioutil.ReadFile(path.Join(certsDir, "certs/ca.cert.pem"))
-	if err != nil {
-		logger.Fatalf("Cannot read CA cert: %v", err)
-	}
-
-	caKeyBytes, err := ioutil.ReadFile(path.Join(certsDir, "private/ca.key.pem"))
-	if err != nil {
-		logger.Fatalf("Cannot read CA key: %v", err)
-	}
-
-	caCertPB, _ := pem.Decode(caCertBytes)
-	caKeyPB, _ := pem.Decode(caKeyBytes)
-	caCert, err := x509.ParseCertificate(caCertPB.Bytes)
-	if err != nil {
-		logger.Fatalf("Cannot parse CA cert: %v", err)
-	}
-
-	caPrivKey, err := x509.ParsePKCS1PrivateKey(caKeyPB.Bytes)
-	if err != nil {
-		logger.Fatalf("Cannot parse CA key: %v", err)
-	}
-
-	// Generate a private key for new client cert.
-	certPrivKey, err := rsa.GenerateKey(rand.Reader, 4096)
-	if err != nil {
-		err := fmt.Errorf("cannot generate private key: %v", err)
-		return nil, err
-	}
-
-	// Prepare certificate template.
-	template := &x509.Certificate{
-		SerialNumber: big.NewInt(1),
-		Subject: pkix.Name{
-			CommonName:    "OpAMP Example Client",
-			Organization:  []string{"OpAMP Example"},
-			Country:       []string{"CA"},
-			Province:      []string{"ON"},
-			Locality:      []string{"City"},
-			StreetAddress: []string{""},
-			PostalCode:    []string{""},
-		},
-		IPAddresses: []net.IP{net.IPv4(127, 0, 0, 1)},
-		NotBefore:   time.Now(),
-		NotAfter:    time.Now().Add(time.Hour * 1000),
-		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
-		KeyUsage:    x509.KeyUsageDigitalSignature,
-	}
-
-	// Create the client cert. Sign it using CA cert.
-	certBytes, err := x509.CreateCertificate(rand.Reader, template, caCert, &certPrivKey.PublicKey, caPrivKey)
-	if err != nil {
-		err := fmt.Errorf("cannot create certificate: %v", err)
-		return nil, err
-	}
-
-	publicKeyPEM := new(bytes.Buffer)
-	pem.Encode(publicKeyPEM, &pem.Block{
-		Type:  "CERTIFICATE",
-		Bytes: certBytes,
-	})
-
-	privateKeyPEM := new(bytes.Buffer)
-	pem.Encode(privateKeyPEM, &pem.Block{
-		Type:  "RSA PRIVATE KEY",
-		Bytes: x509.MarshalPKCS1PrivateKey(certPrivKey),
-	})
-
-	// We have a client certificate with a public and private key.
-	certificate := &protobufs.TLSCertificate{
-		PublicKey:   publicKeyPEM.Bytes(),
-		PrivateKey:  privateKeyPEM.Bytes(),
-		CaPublicKey: caCertBytes,
-	}
-
-	return certificate, nil
-}
diff --git a/internal/tools/go.mod b/internal/tools/go.mod
index 0384b31..2554bd8 100644
--- a/internal/tools/go.mod
+++ b/internal/tools/go.mod
@@ -1,6 +1,6 @@
 module github.com/open-telemetry/opamp-go/internal/tools
 
-go 1.17
+go 1.20
 
 require github.com/ory/go-acc v0.2.8
 
diff --git a/server/serverimpl.go b/server/serverimpl.go
index 5fc7cac..cd19d68 100644
--- a/server/serverimpl.go
+++ b/server/serverimpl.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"sync"
 
 	"github.com/gorilla/websocket"
 	"google.golang.org/protobuf/proto"
@@ -179,7 +180,7 @@ func (s *server) httpHandler(w http.ResponseWriter, req *http.Request) {
 }
 
 func (s *server) handleWSConnection(wsConn *websocket.Conn, connectionCallbacks serverTypes.ConnectionCallbacks) {
-	agentConn := wsConnection{wsConn: wsConn}
+	agentConn := wsConnection{wsConn: wsConn, connMutex: &sync.Mutex{}}
 
 	defer func() {
 		// Close the connection when all is done.
diff --git a/server/wsconnection.go b/server/wsconnection.go
index 9ce53a4..b2b9921 100644
--- a/server/wsconnection.go
+++ b/server/wsconnection.go
@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"net"
+	"sync"
 
 	"github.com/gorilla/websocket"
 
@@ -13,7 +14,11 @@ import (
 
 // wsConnection represents a persistent OpAMP connection over a WebSocket.
 type wsConnection struct {
-	wsConn *websocket.Conn
+	// The websocket library does not allow multiple concurrent write operations,
+	// so ensure that we only have a single operation in progress at a time.
+	// For more: https://pkg.go.dev/github.com/gorilla/websocket#hdr-Concurrency
+	connMutex *sync.Mutex
+	wsConn    *websocket.Conn
 }
 
 var _ types.Connection = (*wsConnection)(nil)
@@ -22,10 +27,10 @@ func (c wsConnection) Connection() net.Conn {
 	return c.wsConn.UnderlyingConn()
 }
 
-// Message header is currently uint64 zero value.
-const wsMsgHeader = uint64(0)
-
 func (c wsConnection) Send(_ context.Context, message *protobufs.ServerToAgent) error {
+	c.connMutex.Lock()
+	defer c.connMutex.Unlock()
+
 	return internal.WriteWSMessage(c.wsConn, message)
 }
 


Test patch:
diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
index 546e056..c45ba20 100644
--- a/.github/workflows/build-and-test.yml
+++ b/.github/workflows/build-and-test.yml
@@ -15,7 +15,7 @@ jobs:
       - name: Setup Go
         uses: actions/setup-go@v2.1.4
         with:
-          go-version: 1.17
+          go-version: "1.20"
       - name: Setup Go Environment
         run: |
           echo "GOPATH=$(go env GOPATH)" >> $GITHUB_ENV
@@ -48,7 +48,7 @@ jobs:
       - name: Setup Go
         uses: actions/setup-go@v2.1.4
         with:
-          go-version: 1.17
+          go-version: "1.20"
       - name: Setup Go Environment
         run: |
           echo "GOPATH=$(go env GOPATH)" >> $GITHUB_ENV
@@ -76,7 +76,7 @@ jobs:
       - name: Setup Go
         uses: actions/setup-go@v3
         with:
-          go-version: 1.17
+          go-version: "1.20"
       - name: Setup Go Environment
         run: |
           echo "GOPATH=$(go env GOPATH)" >> $GITHUB_ENV
diff --git a/server/serverimpl_test.go b/server/serverimpl_test.go
index 7ebaaa3..04827d3 100644
--- a/server/serverimpl_test.go
+++ b/server/serverimpl_test.go
@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -723,3 +724,124 @@ func TestDecodeMessage(t *testing.T) {
 		}
 	}
 }
+
+func TestConnectionAllowsConcurrentWrites(t *testing.T) {
+	srvConnVal := atomic.Value{}
+	callbacks := CallbacksStruct{
+		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnConnectedFunc: func(conn types.Connection) {
+					srvConnVal.Store(conn)
+				},
+			}}
+		},
+	}
+
+	// Start a Server.
+	settings := &StartSettings{Settings: Settings{Callbacks: callbacks}}
+	srv := startServer(t, settings)
+	defer srv.Stop(context.Background())
+
+	// Connect to the Server.
+	conn, _, err := dialClient(settings)
+
+	// Verify that the connection is successful.
+	assert.NoError(t, err)
+	assert.NotNil(t, conn)
+
+	defer conn.Close()
+
+	timeout, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+
+	select {
+	case <-timeout.Done():
+		t.Error("Client failed to connect before timeout")
+	default:
+		if _, ok := srvConnVal.Load().(types.Connection); ok == true {
+			break
+		}
+	}
+
+	cancel()
+
+	srvConn := srvConnVal.Load().(types.Connection)
+	for i := 0; i < 20; i++ {
+		go func() {
+			defer func() {
+				if recover() != nil {
+					require.Fail(t, "Sending to client panicked")
+				}
+			}()
+
+			srvConn.Send(context.Background(), &protobufs.ServerToAgent{})
+		}()
+	}
+}
+
+func BenchmarkSendToClient(b *testing.B) {
+	clientConnections := []*websocket.Conn{}
+	serverConnections := []types.Connection{}
+	srvConnectionsMutex := sync.Mutex{}
+	callbacks := CallbacksStruct{
+		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnConnectedFunc: func(conn types.Connection) {
+					srvConnectionsMutex.Lock()
+					serverConnections = append(serverConnections, conn)
+					srvConnectionsMutex.Unlock()
+				},
+			}}
+		},
+	}
+
+	// Start a Server.
+	settings := &StartSettings{
+		Settings:       Settings{Callbacks: callbacks},
+		ListenEndpoint: testhelpers.GetAvailableLocalAddress(),
+		ListenPath:     "/",
+	}
+	srv := New(&sharedinternal.NopLogger{})
+	err := srv.Start(*settings)
+
+	if err != nil {
+		b.Error(err)
+	}
+
+	defer srv.Stop(context.Background())
+
+	for i := 0; i < b.N; i++ {
+		conn, resp, err := dialClient(settings)
+
+		if err != nil || resp == nil || conn == nil {
+			b.Error("Could not establish connection:", err)
+		}
+
+		clientConnections = append(clientConnections, conn)
+	}
+
+	timeout, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+
+	select {
+	case <-timeout.Done():
+		b.Error("Connections failed to establish in time")
+	default:
+		if len(serverConnections) == b.N {
+			break
+		}
+	}
+
+	cancel()
+
+	for _, conn := range serverConnections {
+		err := conn.Send(context.Background(), &protobufs.ServerToAgent{})
+
+		if err != nil {
+			b.Error(err)
+		}
+	}
+
+	for _, conn := range clientConnections {
+		conn.Close()
+	}
+
+}