Use ClientOpt and ServerOpt for websocket TLS config

Signed-off-by: Lorenzo <[email protected]>

Author: lorenzodonini

example/1.6/cp/charge_point_sim.go

@@ -64,10 +64,10 @@ func setupTlsChargePoint(chargePointID string) ocpp16.ChargePoint {
 		}
 	}
 	// Create client with TLS config
-	client := ws.NewTLSClient(&tls.Config{
+	client := ws.NewClient(ws.WithClientTLSConfig(&tls.Config{
 		RootCAs:      certPool,
 		Certificates: clientCertificates,
-	})
+	}))
 	return ocpp16.NewChargePoint(chargePointID, nil, client)
 }
 

example/1.6/cs/central_system_sim.go

@@ -67,10 +67,10 @@ func setupTlsCentralSystem() ocpp16.CentralSystem {
 	if !ok {
 		log.Fatalf("no required %v found", envVarServerCertificateKey)
 	}
-	server := ws.NewTLSServer(certificate, key, &tls.Config{
+	server := ws.NewServer(ws.WithServerTLSConfig(certificate, key, &tls.Config{
 		ClientAuth: tls.RequireAndVerifyClientCert,
 		ClientCAs:  certPool,
-	})
+	}))
 	return ocpp16.NewCentralSystem(nil, server)
 }
 

example/2.0.1/chargingstation/charging_station_sim.go

@@ -66,10 +66,10 @@ func setupTlsChargingStation(chargingStationID string) ocpp2.ChargingStation {
 		}
 	}
 	// Create client with TLS config
-	client := ws.NewTLSClient(&tls.Config{
+	client := ws.NewClient(ws.WithClientTLSConfig(&tls.Config{
 		RootCAs:      certPool,
 		Certificates: clientCertificates,
-	})
+	}))
 	return ocpp2.NewChargingStation(chargingStationID, nil, client)
 }
 

example/2.0.1/csms/csms_sim.go

@@ -70,10 +70,10 @@ func setupTlsCentralSystem() ocpp2.CSMS {
 	if !ok {
 		log.Fatalf("no required %v found", envVarServerCertificateKey)
 	}
-	server := ws.NewTLSServer(certificate, key, &tls.Config{
+	server := ws.NewServer(ws.WithServerTLSConfig(certificate, key, &tls.Config{
 		ClientAuth: tls.RequireAndVerifyClientCert,
 		ClientCAs:  certPool,
-	})
+	}))
 	return ocpp2.NewCSMS(nil, server)
 }
 

ocpp1.6/v16.go

@@ -153,9 +153,9 @@ type ChargePoint interface {
 //	if !ok {
 //		log.Fatal("couldn't parse PEM certificate")
 //	}
-//	cp := NewClient("someUniqueId", nil, ws.NewTLSClient(&tls.Config{
+//	cp := NewClient("someUniqueId", nil, ws.NewClient(ws.WithClientTLSConfig(&tls.Config{
 //		RootCAs: certPool,
-//	})
+//	}))
 //
 // For more advanced options, or if a customer networking/occpj layer is required,
 // please refer to ocppj.Client and ws.Client.
@@ -338,7 +338,7 @@ type CentralSystem interface {
 //
 // If you need a TLS server, you may use the following:
 //
-//	cs := NewServer(nil, ws.NewTLSServer("certificatePath", "privateKeyPath"))
+//	cs := NewServer(nil, ws.NewServer(ws.WithServerTLSConfig("certificatePath", "privateKeyPath", nil)))
 func NewCentralSystem(endpoint *ocppj.Server, server ws.Server) CentralSystem {
 	if server == nil {
 		server = ws.NewServer()

ocpp2.0.1/v2.go

@@ -202,9 +202,9 @@ type ChargingStation interface {
 //	if !ok {
 //		log.Fatal("couldn't parse PEM certificate")
 //	}
-//	cs := NewChargingStation("someUniqueId", nil, ws.NewTLSClient(&tls.Config{
+//	cs := NewChargingStation("someUniqueId", nil, ws.NewClient(ws.WithClientTLSConfig(&tls.Config{
 //		RootCAs: certPool,
-//	})
+//	}))
 //
 // For more advanced options, or if a custom networking/occpj layer is required,
 // please refer to ocppj.Client and ws.Client.
@@ -414,7 +414,7 @@ type CSMS interface {
 //
 // If you need a TLS server, you may use the following:
 //
-//	csms := NewCSMS(nil, ws.NewTLSServer("certificatePath", "privateKeyPath"))
+//	csms := NewCSMS(nil, ws.NewServer(ws.WithServerTLSConfig("certificatePath", "privateKeyPath", nil)))
 func NewCSMS(endpoint *ocppj.Server, server ws.Server) CSMS {
 	if server == nil {
 		server = ws.NewServer()

ws/client.go

@@ -23,14 +23,14 @@ import (
 //
 //	client := NewClient()
 //
-// If you need a TLS ws client instead, use:
+// If you need a secure websocket client instead, pass a tls.Config to the NewClient function:
 //
 //	certPool, err := x509.SystemCertPool()
 //	if err != nil {
 //		log.Fatal(err)
 //	}
-//	// You may add more trusted certificates to the pool before creating the TLSClientConfig
-//	client := NewTLSClient(&tls.Config{
+//	// You may add more trusted certificates to the pool before creating the TLS Config
+//	client := NewClient(&tls.Config{
 //		RootCAs: certPool,
 //	})
 //
@@ -122,7 +122,7 @@ type Client interface {
 
 // client is the default implementation of a Websocket client.
 //
-// Use the NewClient or NewTLSClient functions to create a new client.
+// Use the NewClient function to create a new client.
 type client struct {
 	webSocket      *webSocket
 	url            url.URL
@@ -136,52 +136,57 @@ type client struct {
 	reconnectC     chan struct{} // used for signaling, that a reconnection attempt should be interrupted
 }
 
-// Creates a new simple websocket client (the channel is not secured).
-//
-// Additional options may be added using the AddOption function.
-//
-// Basic authentication can be set using the SetBasicAuth function.
-//
-// By default, the client will not neogtiate any subprotocol. This value needs to be set via the
-// respective SetRequestedSubProtocol method.
-func NewClient() Client {
-	return &client{
-		dialOptions:   []func(*websocket.Dialer){},
-		timeoutConfig: NewClientTimeoutConfig(),
-		reconnectC:    make(chan struct{}, 1),
-		header:        http.Header{},
+// ClientOpt is a function that can be used to set options on a client during creation.
+type ClientOpt func(c *client)
+
+// WithClientTLSConfig sets the TLS configuration for the client.
+// If the passed tlsConfig is nil, the client will not use TLS.
+func WithClientTLSConfig(tlsConfig *tls.Config) ClientOpt {
+	return func(c *client) {
+		if tlsConfig != nil {
+			c.dialOptions = append(c.dialOptions, func(dialer *websocket.Dialer) {
+				dialer.TLSClientConfig = tlsConfig
+			})
+		}
 	}
 }
 
-// NewTLSClient creates a new secure websocket client. If supported by the server, the websocket channel will use TLS.
+// NewClient creates a new websocket client.
+//
+// If the optional tlsConfig is not nil, and the server supports secure communication,
+// the websocket channel will use TLS.
 //
 // Additional options may be added using the AddOption function.
+//
 // Basic authentication can be set using the SetBasicAuth function.
 //
+// By default, the client will not negotiate any sub-protocol. This value needs to be set via the
+// respective SetRequestedSubProtocol method.
+//
 // To set a client certificate, you may do:
 //
 //	certificate, _ := tls.LoadX509KeyPair(clientCertPath, clientKeyPath)
 //	clientCertificates := []tls.Certificate{certificate}
-//	client := ws.NewTLSClient(&tls.Config{
+//	client := ws.NewClient(ws.WithClientTLSConfig(&tls.Config{
 //		RootCAs:      certPool,
 //		Certificates: clientCertificates,
-//	})
+//	}))
 //
-// You can set any other TLS option within the same constructor as well.
+// You can set any other TLS option within the same constructor config as well.
 // For example, if you wish to test connecting to a server having a
 // self-signed certificate (do not use in production!), pass:
 //
 //	InsecureSkipVerify: true
-func NewTLSClient(tlsConfig *tls.Config) Client {
+func NewClient(opts ...ClientOpt) Client {
 	c := &client{
 		dialOptions:   []func(*websocket.Dialer){},
 		timeoutConfig: NewClientTimeoutConfig(),
 		reconnectC:    make(chan struct{}, 1),
 		header:        http.Header{},
 	}
-	c.dialOptions = append(c.dialOptions, func(dialer *websocket.Dialer) {
-		dialer.TLSClientConfig = tlsConfig
-	})
+	for _, o := range opts {
+		o(c)
+	}
 	return c
 }
 

ws/server.go

@@ -26,9 +26,9 @@ type CheckClientHandler func(id string, r *http.Request) bool
 //
 //	server := NewServer()
 //
-// If you need a TLS ws server instead, use:
+// If you need a server with TLS support, pass the following option:
 //
-//	server := NewTLSServer("cert.pem", "privateKey.pem")
+//	server := NewServer(WithServerTLSConfig("cert.pem", "privateKey.pem", nil))
 //
 // To support client basic authentication, use:
 //
@@ -125,7 +125,7 @@ type Server interface {
 
 // Default implementation of a Websocket server.
 //
-// Use the NewServer or NewTLSServer functions to create a new server.
+// Use the NewServer function to create a new server.
 type server struct {
 	connections         map[string]*webSocket
 	httpServer          *http.Server
@@ -144,43 +144,50 @@ type server struct {
 	httpHandler         *mux.Router
 }
 
-// Creates a new simple websocket server (the websockets are not secured).
-func NewServer() Server {
-	router := mux.NewRouter()
-	return &server{
-		httpServer:    &http.Server{},
-		timeoutConfig: NewServerTimeoutConfig(),
-		upgrader:      websocket.Upgrader{Subprotocols: []string{}},
-		httpHandler:   router,
+// ServerOpt is a function that can be used to set options on a server during creation.
+type ServerOpt func(s *server)
+
+// WithServerTLSConfig sets the TLS configuration for the server.
+// If the passed tlsConfig is nil, the client will not use TLS.
+func WithServerTLSConfig(certificatePath string, certificateKey string, tlsConfig *tls.Config) ServerOpt {
+	return func(s *server) {
+		s.tlsCertificatePath = certificatePath
+		s.tlsCertificateKey = certificateKey
+		if tlsConfig != nil {
+			s.httpServer.TLSConfig = tlsConfig
+		}
 	}
 }
 
-// NewTLSServer creates a new secure websocket server. All created websocket channels will use TLS.
+// NewServer Creates a new websocket server.
+//
+// Additional options may be added using the AddOption function.
 //
-// You need to pass a filepath to the server TLS certificate and key.
+// By default, the websockets are not secure, and the server will not perform any client certificate verification.
 //
-// It is recommended to pass a valid TLSConfig for the server to use.
-// For example to require client certificate verification:
+// To add TLS support to the server, a valid server certificate path and key must be passed.
+// To also add support for client certificate verification, a valid TLSConfig needs to be configured.
+// For example:
 //
-//	tlsConfig := &tls.Config{
-//		ClientAuth: tls.RequireAndVerifyClientCert,
-//		ClientCAs: clientCAs,
-//	}
+//		tlsConfig := &tls.Config{
+//			ClientAuth: tls.RequireAndVerifyClientCert,
+//			ClientCAs: clientCAs,
+//		}
+//	 server := ws.NewServer(ws.WithServerTLSConfig("cert.pem", "privateKey.pem", tlsConfig))
 //
-// If no tlsConfig parameter is passed, the server will by default
-// not perform any client certificate verification.
-func NewTLSServer(certificatePath string, certificateKey string, tlsConfig *tls.Config) Server {
+// When TLS is correctly configured, the server will automatically use it for all created websocket channels.
+func NewServer(opts ...ServerOpt) Server {
 	router := mux.NewRouter()
-	return &server{
-		tlsCertificatePath: certificatePath,
-		tlsCertificateKey:  certificateKey,
-		httpServer: &http.Server{
-			TLSConfig: tlsConfig,
-		},
+	s := &server{
+		httpServer:    &http.Server{},
 		timeoutConfig: NewServerTimeoutConfig(),
 		upgrader:      websocket.Upgrader{Subprotocols: []string{}},
 		httpHandler:   router,
 	}
+	for _, o := range opts {
+		o(s)
+	}
+	return s
 }
 
 func (s *server) SetMessageHandler(handler MessageHandler) {