From 8c07e71d01dbc8cb507648fc3a625065e40d0bb3 Mon Sep 17 00:00:00 2001
From: Siyuan Miao <i@xswan.net>
Date: Fri, 7 Mar 2025 20:14:22 +0100
Subject: [PATCH] feat(tls): add simple tls support

---
 config.go  |   2 +
 main.go    |   3 ++
 web_tls.go | 132 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 137 insertions(+)
 create mode 100644 web_tls.go

diff --git a/config.go b/config.go
index e4e27d7..304d8f6 100644
--- a/config.go
+++ b/config.go
@@ -37,6 +37,7 @@ type Config struct {
 	DisplayMaxBrightness int               `json:"display_max_brightness"`
 	DisplayDimAfterSec   int               `json:"display_dim_after_sec"`
 	DisplayOffAfterSec   int               `json:"display_off_after_sec"`
+	TLSMode              string            `json:"tls_mode"`
 	UsbConfig            *UsbConfig        `json:"usb_config"`
 }
 
@@ -50,6 +51,7 @@ var defaultConfig = &Config{
 	DisplayMaxBrightness: 64,
 	DisplayDimAfterSec:   120,  // 2 minutes
 	DisplayOffAfterSec:   1800, // 30 minutes
+	TLSMode:              "",
 	UsbConfig: &UsbConfig{
 		VendorId:     "0x1d6b", //The Linux Foundation
 		ProductId:    "0x0104", //Multifunction Composite Gadget
diff --git a/main.go b/main.go
index a3caba7..2c8c22a 100644
--- a/main.go
+++ b/main.go
@@ -68,6 +68,9 @@ func Main() {
 	}()
 	//go RunFuseServer()
 	go RunWebServer()
+	if config.TLSMode != "" {
+		go RunWebSecureServer()
+	}
 	// If the cloud token isn't set, the client won't be started by default.
 	// However, if the user adopts the device via the web interface, handleCloudRegister will start the client.
 	if config.CloudToken != "" {
diff --git a/web_tls.go b/web_tls.go
new file mode 100644
index 0000000..fff9253
--- /dev/null
+++ b/web_tls.go
@@ -0,0 +1,132 @@
+package kvm
+
+import (
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/tls"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"encoding/pem"
+	"log"
+	"math/big"
+	"net"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+)
+
+const (
+	WebSecureListen                  = ":443"
+	WebSecureSelfSignedDefaultDomain = "jetkvm.local"
+	WebSecureSelfSignedDuration      = 365 * 24 * time.Hour
+)
+
+var (
+	tlsCerts    = make(map[string]*tls.Certificate)
+	tlsCertLock = &sync.Mutex{}
+)
+
+// RunWebSecureServer runs a web server with TLS.
+func RunWebSecureServer() {
+	r := setupRouter()
+
+	server := &http.Server{
+		Addr:    WebSecureListen,
+		Handler: r,
+		TLSConfig: &tls.Config{
+			// TODO: cache certificate in persistent storage
+			GetCertificate: func(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
+				hostname := WebSecureSelfSignedDefaultDomain
+				if info.ServerName != "" {
+					hostname = info.ServerName
+				} else {
+					hostname = strings.Split(info.Conn.LocalAddr().String(), ":")[0]
+				}
+
+				logger.Infof("TLS handshake for %s, SupportedProtos: %v", hostname, info.SupportedProtos)
+
+				cert := createSelfSignedCert(hostname)
+
+				return cert, nil
+			},
+		},
+	}
+	logger.Infof("Starting websecure server on %s", RunWebSecureServer)
+	err := server.ListenAndServeTLS("", "")
+	if err != nil {
+		panic(err)
+	}
+	return
+}
+
+func createSelfSignedCert(hostname string) *tls.Certificate {
+	if tlsCert := tlsCerts[hostname]; tlsCert != nil {
+		return tlsCert
+	}
+	tlsCertLock.Lock()
+	defer tlsCertLock.Unlock()
+
+	logger.Infof("Creating self-signed certificate for %s", hostname)
+
+	priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	if err != nil {
+		log.Fatalf("Failed to generate private key: %v", err)
+	}
+	keyUsage := x509.KeyUsageDigitalSignature
+
+	notBefore := time.Now()
+	notAfter := notBefore.AddDate(1, 0, 0)
+
+	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
+	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
+	if err != nil {
+		logger.Errorf("Failed to generate serial number: %v", err)
+	}
+
+	dnsName := hostname
+	ip := net.ParseIP(hostname)
+	if ip != nil {
+		dnsName = WebSecureSelfSignedDefaultDomain
+	}
+
+	template := x509.Certificate{
+		SerialNumber: serialNumber,
+		Subject: pkix.Name{
+			CommonName:   hostname,
+			Organization: []string{"JetKVM"},
+		},
+		NotBefore: notBefore,
+		NotAfter:  notAfter,
+
+		KeyUsage:              keyUsage,
+		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
+		BasicConstraintsValid: true,
+
+		DNSNames:    []string{dnsName},
+		IPAddresses: []net.IP{},
+	}
+
+	if ip != nil {
+		template.IPAddresses = append(template.IPAddresses, ip)
+	}
+
+	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
+	if err != nil {
+		logger.Errorf("Failed to create certificate: %v", err)
+	}
+
+	cert := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
+	if cert == nil {
+		logger.Errorf("Failed to encode certificate")
+	}
+
+	tlsCert := &tls.Certificate{
+		Certificate: [][]byte{derBytes},
+		PrivateKey:  priv,
+	}
+	tlsCerts[hostname] = tlsCert
+
+	return tlsCert
+}