From 2f9cb7377bee4cde0f5ae018cccba0f1fa33b705 Mon Sep 17 00:00:00 2001
From: Andy Ye <35905412+TurboTheTurtle@users.noreply.github.com>
Date: Tue, 12 May 2026 12:24:39 -0700
Subject: [PATCH] Pin client TLS certificate

---
 main.go      | 30 ++++++++++++++++++++++++++++++
 main_test.go | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)
 create mode 100644 main_test.go

diff --git a/main.go b/main.go
index 51d1970..737d0d1 100644
--- a/main.go
+++ b/main.go
@@ -1,6 +1,8 @@
 package main
 
 import (
+	"crypto/sha256"
+	"encoding/pem"
 	"flag"
 	"fmt"
 	"os"
@@ -80,6 +82,28 @@ func readCertificate() ([]byte, error) {
 	panic("thou shalt not reach hear")
 }
 
+func calculatePEMCertChainHash(certContent []byte) ([]byte, error) {
+	var hashValue []byte
+	for {
+		block, remain := pem.Decode(certContent)
+		if block == nil {
+			break
+		}
+		out := sha256.Sum256(block.Bytes)
+		if hashValue == nil {
+			hashValue = out[:]
+		} else {
+			newHashValue := sha256.Sum256(append(hashValue, out[:]...))
+			hashValue = newHashValue[:]
+		}
+		certContent = remain
+	}
+	if hashValue == nil {
+		return nil, newError("no certificate found")
+	}
+	return hashValue, nil
+}
+
 func logConfig(logLevel string) *vlog.Config {
 	config := &vlog.Config{
 		Error:  &vlog.LogSpecification{Type: vlog.LogType_Console, Level: clog.Severity_Warning},
@@ -189,6 +213,12 @@ func generateConfig() (*core.Config, error) {
 			if err != nil {
 				return nil, newError("failed to read cert").Base(err)
 			}
+			certHash, err := calculatePEMCertChainHash(certificate.Certificate)
+			if err != nil {
+				return nil, newError("failed to parse cert").Base(err)
+			}
+			tlsConfig.AllowInsecureIfPinnedPeerCertificate = true
+			tlsConfig.PinnedPeerCertificateChainSha256 = [][]byte{certHash}
 			tlsConfig.Certificate = []*tls.Certificate{&certificate}
 		}
 		streamConfig.SecurityType = serial.GetMessageType(&tlsConfig)
diff --git a/main_test.go b/main_test.go
new file mode 100644
index 0000000..b8c9ffa
--- /dev/null
+++ b/main_test.go
@@ -0,0 +1,35 @@
+package main
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"encoding/pem"
+	"testing"
+)
+
+func TestCalculatePEMCertChainHash(t *testing.T) {
+	leaf := []byte("leaf certificate")
+	issuer := []byte("issuer certificate")
+	pemChain := bytes.Join([][]byte{
+		pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: leaf}),
+		pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: issuer}),
+	}, nil)
+
+	hash, err := calculatePEMCertChainHash(pemChain)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	leafHash := sha256.Sum256(leaf)
+	issuerHash := sha256.Sum256(issuer)
+	expected := sha256.Sum256(append(leafHash[:], issuerHash[:]...))
+	if !bytes.Equal(hash, expected[:]) {
+		t.Fatalf("unexpected hash: got %x, want %x", hash, expected)
+	}
+}
+
+func TestCalculatePEMCertChainHashRejectsInvalidPEM(t *testing.T) {
+	if _, err := calculatePEMCertChainHash([]byte("not a certificate")); err == nil {
+		t.Fatal("expected invalid PEM to fail")
+	}
+}