Add PSK support

client/client.go

@@ -4,6 +4,7 @@
 // Use ClientConfig struct to set various client's options and features.
 //
 // Example:
+//
 //	import (
 //		"fmt"
 //		"time"
@@ -58,12 +59,16 @@ import (
 	"crypto"
 	"crypto/tls"
 	"crypto/x509"
+	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net"
+	"net/http"
 	"sync"
 	"time"
 
+	"github.com/patrickmn/go-cache"
 	"github.com/pion/dtls/v2"
 	"github.com/pion/dtls/v2/pkg/crypto/selfsign"
 	"golang.org/x/sync/errgroup"
@@ -78,10 +83,23 @@ import (
 type ClientConfig struct {
 	// UseDTLS controls whether DTLS should be used to secure the connection
 	// to the MQTT-SN gateway.
-	UseDTLS        bool
-	Certificate    *tls.Certificate
-	PrivateKey     crypto.PrivateKey
-	CACertificates []*x509.Certificate
+	UseDTLS     bool
+	Certificate *tls.Certificate
+	PrivateKey  crypto.PrivateKey
+	// UsePSK controls whether pre-shared key should be used to secure the
+	// connection to the MQTT-SN gateway. If UsePSK is true, you must provide
+	// PSKIdentity, PSKApiBasicAuthUsername, PSKApiBasicAuthPassword and
+	// PSKApiEndpoint.
+	// If UsePSK is true, the client will use PSK instead of the certificate
+	// and private key.
+	UsePSK                  bool
+	PSK                     *cache.Cache
+	PSKCacheExpiration      time.Duration
+	PSKIdentity             string
+	PSKApiBasicAuthUsername string
+	PSKApiBasicAuthPassword string
+	PSKApiEndpoint          string
+	CACertificates          []*x509.Certificate
 	// SelfSigned controls whether the client should use a self-signed
 	// certificate and key.  If SelfSigned is false and UseDTLS is true, you
 	// must provide CertFile and KeyFile.
@@ -147,21 +165,24 @@ func (c *Client) connectDTLS(ctx context.Context, address string) (net.Conn, err
 	var certificate *tls.Certificate
 	var err error
 
-	if c.cfg.SelfSigned {
-		var cert tls.Certificate
-		cert, err = selfsign.GenerateSelfSigned()
-		certificate = &cert
-	} else {
-		privateKey := c.cfg.PrivateKey
-		if privateKey == nil {
-			err = errors.New("private key is missing")
-		}
-		if certificate = c.cfg.Certificate; certificate != nil {
-			certificate.PrivateKey = privateKey
+	if !c.cfg.UsePSK && c.cfg.UseDTLS {
+		if c.cfg.SelfSigned {
+			var cert tls.Certificate
+			cert, err = selfsign.GenerateSelfSigned()
+			certificate = &cert
 		} else {
-			err = errors.New("TLS certificate is missing")
+			privateKey := c.cfg.PrivateKey
+			if privateKey == nil {
+				err = errors.New("private key is missing")
+			}
+			if certificate = c.cfg.Certificate; certificate != nil {
+				certificate.PrivateKey = privateKey
+			} else {
+				err = errors.New("TLS certificate is missing")
+			}
 		}
 	}
+
 	if err != nil {
 		return nil, err
 	}
@@ -182,12 +203,34 @@ func (c *Client) connectDTLS(ctx context.Context, address string) (net.Conn, err
 
 	// Prepare the configuration of the DTLS connection
 	config := &dtls.Config{
-		Certificates:         []tls.Certificate{*certificate},
 		InsecureSkipVerify:   c.cfg.Insecure,
 		ExtendedMasterSecret: dtls.RequireExtendedMasterSecret,
 		RootCAs:              certPool,
 	}
 
+	if !c.cfg.UsePSK && c.cfg.UseDTLS && certificate != nil {
+		config.Certificates = []tls.Certificate{*certificate}
+	}
+
+	if c.cfg.UsePSK && c.cfg.UseDTLS {
+		config.CipherSuites = []dtls.CipherSuiteID{dtls.TLS_PSK_WITH_AES_128_GCM_SHA256}
+		config.PSK = func(hint []byte) ([]byte, error) {
+			psk, ok := c.cfg.PSK.Get(string(hint))
+			if ok {
+				return psk.([]byte), nil
+			}
+
+			psk, ok = c.getPSK(string(hint))
+			if ok {
+				c.cfg.PSK.Set(string(hint), psk, c.cfg.PSKCacheExpiration)
+				return psk.([]byte), nil
+			}
+
+			return nil, errors.New("PSK not found")
+		}
+		config.PSKIdentityHint = []byte(c.cfg.PSKIdentity)
+	}
+
 	// Connect to a DTLS server
 	ctx2, cancel := context.WithTimeout(ctx, 30*time.Second)
 	defer cancel()
@@ -556,3 +599,47 @@ func (c *Client) Disconnect() error {
 		return c.group.Wait()
 	}
 }
+
+func (c *Client) getPSK(hint string) ([]byte, bool) {
+	req, err := http.NewRequest("GET", fmt.Sprintf(c.cfg.PSKApiEndpoint+"/%s", hint), nil)
+	if err != nil {
+		c.log.Error("Error in creating request: %s", err)
+		return nil, false
+	}
+
+	req.SetBasicAuth(c.cfg.PSKApiBasicAuthUsername, c.cfg.PSKApiBasicAuthPassword)
+	client := &http.Client{}
+	resp, err := client.Do(req)
+
+	if err != nil {
+		c.log.Error("Error in sending request: %s", err)
+		return nil, false
+	}
+
+	defer resp.Body.Close()
+
+	if resp.StatusCode == http.StatusNotFound {
+		c.log.Debug("ID not found")
+		return nil, false
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		c.log.Error("Error in reading response body: %s", err)
+		return nil, false
+	}
+
+	type Response struct {
+		Data map[string][]byte `json:"data"`
+	}
+
+	var response Response
+
+	err = json.Unmarshal(body, &response)
+	if err != nil {
+		c.log.Error("Error in unmarshalling response body: %s", err)
+		return nil, false
+	}
+
+	return response.Data[hint], true
+}

cmd/bisquitt-pub/actions.go

@@ -13,6 +13,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/patrickmn/go-cache"
 	"github.com/urfave/cli/v2"
 
 	snClient "github.com/energomonitor/bisquitt/client"
@@ -33,13 +34,20 @@ func handleAction() cli.ActionFunc {
 
 		useDTLS := c.Bool(DtlsFlag)
 		useSelfSigned := c.Bool(SelfSignedFlag)
+		usePSK := c.Bool(PskFlag)
+		pskCacheExpiration := c.Duration(PskCacheExpirationFlag)
+
+		pskIdentity := c.String(PskIdentityFlag)
+		pskApiBasicAuthUsername := c.String(PskApiBasicAuthUsernameFlag)
+		pskApiBasicAuthPassword := c.String(PskApiBasicAuthPasswordFlag)
+		pskApiEndpoint := c.String(PskApiEndpointFlag)
 		certFile := c.Path(CertFlag)
 		keyFile := c.Path(KeyFlag)
 		caFile := c.Path(CAFileFlag)
 		caPath := c.Path(CAPathFlag)
 		debug := c.Bool(DebugFlag)
 
-		if useDTLS && (certFile == "" || keyFile == "") && !useSelfSigned {
+		if useDTLS && ((certFile == "" || keyFile == "") && !useSelfSigned) && !usePSK {
 			return fmt.Errorf(`options "--%s" and "--%s" are mandatory when using DTLS. Use "--%s" to generate self-signed certificate.`,
 				CertFlag, KeyFlag, SelfSignedFlag)
 		}
@@ -132,21 +140,28 @@ func handleAction() cli.ActionFunc {
 		password := []byte(c.String(PasswordFlag))
 
 		clientCfg := &snClient.ClientConfig{
-			ClientID:         clientID,
-			UseDTLS:          useDTLS,
-			SelfSigned:       useSelfSigned,
-			Insecure:         insecure,
-			Certificate:      certificate,
-			PrivateKey:       privateKey,
-			CACertificates:   caCertificates,
-			RetryDelay:       10 * time.Second,
-			RetryCount:       4,
-			ConnectTimeout:   20 * time.Second,
-			KeepAlive:        60 * time.Second,
-			CleanSession:     true,
-			PredefinedTopics: predefinedTopics,
-			User:             user,
-			Password:         password,
+			ClientID:                clientID,
+			UseDTLS:                 useDTLS,
+			UsePSK:                  usePSK,
+			PSK:                     cache.New(pskCacheExpiration, 5*time.Minute),
+			PSKCacheExpiration:      pskCacheExpiration,
+			PSKIdentity:             pskIdentity,
+			PSKApiBasicAuthUsername: pskApiBasicAuthUsername,
+			PSKApiBasicAuthPassword: pskApiBasicAuthPassword,
+			PSKApiEndpoint:          pskApiEndpoint,
+			SelfSigned:              useSelfSigned,
+			Insecure:                insecure,
+			Certificate:             certificate,
+			PrivateKey:              privateKey,
+			CACertificates:          caCertificates,
+			RetryDelay:              10 * time.Second,
+			RetryCount:              4,
+			ConnectTimeout:          20 * time.Second,
+			KeepAlive:               60 * time.Second,
+			CleanSession:            true,
+			PredefinedTopics:        predefinedTopics,
+			User:                    user,
+			Password:                password,
 		}
 
 		var logger util.Logger

cmd/bisquitt-pub/application.go

@@ -2,32 +2,39 @@ package main
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/urfave/cli/v2"
 
 	"github.com/energomonitor/bisquitt"
 )
 
 const (
-	HostFlag                 = "host"
-	PortFlag                 = "port"
-	DtlsFlag                 = "dtls"
-	SelfSignedFlag           = "self-signed"
-	CertFlag                 = "cert"
-	KeyFlag                  = "key"
-	CAFileFlag               = "cafile"
-	CAPathFlag               = "capath"
-	InsecureFlag             = "insecure"
-	DebugFlag                = "debug"
-	TopicFlag                = "topic"
-	MessageFlag              = "message"
-	RetainFlag               = "retain"
-	PredefinedTopicFlag      = "predefined-topic"
-	PredefinedTopicsFileFlag = "predefined-topics-file"
-	QOSFlag                  = "qos"
-	ClientIDFlag             = "client-id"
-	UserFlag                 = "user"
-	PasswordFlag             = "password"
+	HostFlag                    = "host"
+	PortFlag                    = "port"
+	DtlsFlag                    = "dtls"
+	SelfSignedFlag              = "self-signed"
+	PskFlag                     = "psk"
+	PskCacheExpirationFlag      = "psk-cache-expiration"
+	PskIdentityFlag             = "psk-identity"
+	PskApiBasicAuthUsernameFlag = "psk-api-basic-auth-username"
+	PskApiBasicAuthPasswordFlag = "psk-api-basic-auth-password"
+	PskApiEndpointFlag          = "psk-api-endpoint"
+	CertFlag                    = "cert"
+	KeyFlag                     = "key"
+	CAFileFlag                  = "cafile"
+	CAPathFlag                  = "capath"
+	InsecureFlag                = "insecure"
+	DebugFlag                   = "debug"
+	TopicFlag                   = "topic"
+	MessageFlag                 = "message"
+	RetainFlag                  = "retain"
+	PredefinedTopicFlag         = "predefined-topic"
+	PredefinedTopicsFileFlag    = "predefined-topics-file"
+	QOSFlag                     = "qos"
+	ClientIDFlag                = "client-id"
+	UserFlag                    = "user"
+	PasswordFlag                = "password"
 )
 
 func init() {
@@ -76,6 +83,49 @@ var Application = cli.App{
 				"SELF_SIGNED",
 			},
 		},
+		&cli.BoolFlag{
+			Name:  PskFlag,
+			Usage: "use PSK",
+			EnvVars: []string{
+				"PSK_ENABLED",
+			},
+		},
+		&cli.DurationFlag{
+			Name:  PskCacheExpirationFlag,
+			Usage: "PSK cache expiration",
+			Value: 5 * time.Minute,
+			EnvVars: []string{
+				"PSK_CACHE_EXPIRATION",
+			},
+		},
+		&cli.StringFlag{
+			Name:  PskIdentityFlag,
+			Usage: "PSK identity",
+			EnvVars: []string{
+				"PSK_IDENTITY",
+			},
+		},
+		&cli.StringFlag{
+			Name:  PskApiBasicAuthUsernameFlag,
+			Usage: "PSK API basic auth username",
+			EnvVars: []string{
+				"PSK_API_BASIC_AUTH_USERNAME",
+			},
+		},
+		&cli.StringFlag{
+			Name:  PskApiBasicAuthPasswordFlag,
+			Usage: "PSK API basic auth password",
+			EnvVars: []string{
+				"PSK_API_BASIC_AUTH_PASSWORD",
+			},
+		},
+		&cli.StringFlag{
+			Name:  PskApiEndpointFlag,
+			Usage: "PSK API endpoint",
+			EnvVars: []string{
+				"PSK_API_ENDPOINT",
+			},
+		},
 		&cli.PathFlag{
 			Name:  CertFlag,
 			Usage: "DTLS certificate file",