From df6f0b184c5689eab0a4b4be4769902200f2b241 Mon Sep 17 00:00:00 2001
From: Sergey Matveev <stargrave@stargrave.org>
Date: Tue, 6 Sep 2022 12:55:47 +0300
Subject: [PATCH] Compatibility with raw IPv6 addresses as hostname

---
 tls.go          |  5 +++--
 tls/dane.go     | 14 +++++++-------
 tls/dial.go     |  3 +--
 tls/hostport.go | 20 ++++++++++++++++++++
 4 files changed, 31 insertions(+), 11 deletions(-)
 create mode 100644 tls/hostport.go

diff --git a/tls.go b/tls.go
index 05d21db..0fb4ae1 100644
--- a/tls.go
+++ b/tls.go
@@ -25,8 +25,9 @@ import (
 	"fmt"
 	"log"
 	"net/http"
-	"strings"
 	"time"
+
+	ttls "go.stargrave.org/tofuproxy/tls"
 )
 
 var (
@@ -57,7 +58,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 		req.Proto,
 		http.StatusOK, http.StatusText(http.StatusOK),
 	)))
-	host := strings.Split(req.Host, ":")[0]
+	host, _, _ := ttls.SplitHostPort(req.Host)
 	hostCertsM.Lock()
 	keypair, ok := hostCerts[host]
 	if !ok || !keypair.cert.NotAfter.After(time.Now().Add(time.Hour)) {
diff --git a/tls/dane.go b/tls/dane.go
index ba6e7d0..b970cc2 100644
--- a/tls/dane.go
+++ b/tls/dane.go
@@ -25,7 +25,6 @@ import (
 	"encoding/hex"
 	"fmt"
 	"log"
-	"strings"
 
 	"github.com/miekg/dns"
 )
@@ -36,12 +35,13 @@ func DANE(addr string, cert *x509.Certificate) (bool, bool) {
 	if DNSSrv == "" {
 		return false, false
 	}
-	host := addr
-	port := "443"
-	cols := strings.Split(addr, ":")
-	if len(cols) > 1 {
-		host = cols[0]
-		port = cols[1]
+	host, port, err := SplitHostPort(addr)
+	if err != nil {
+		log.Printf("can not split host+port: %s: %+v\n", addr, err)
+		return false, false
+	}
+	if port == "" {
+		port = "443"
 	}
 	m := new(dns.Msg)
 	m.SetQuestion(dns.Fqdn(fmt.Sprintf("_%s._tcp.%s", port, host)), dns.TypeTLSA)
diff --git a/tls/dial.go b/tls/dial.go
index f2286d7..2b61dac 100644
--- a/tls/dial.go
+++ b/tls/dial.go
@@ -24,7 +24,6 @@ import (
 	"crypto/x509"
 	"fmt"
 	"net"
-	"strings"
 
 	"go.cypherpunks.ru/ucspi"
 	"go.stargrave.org/tofuproxy/fifos"
@@ -33,7 +32,7 @@ import (
 var sessionCache = tls.NewLRUClientSessionCache(1024)
 
 func DialTLS(ctx context.Context, network, addr string) (net.Conn, error) {
-	host := strings.Split(addr, ":")[0]
+	host, _, _ := SplitHostPort(addr)
 	ccg := ClientCertificateGetter{host: host}
 	cfg := tls.Config{
 		VerifyPeerCertificate: func(
diff --git a/tls/hostport.go b/tls/hostport.go
new file mode 100644
index 0000000..14b1ce3
--- /dev/null
+++ b/tls/hostport.go
@@ -0,0 +1,20 @@
+package tofuproxy
+
+import (
+	"net"
+	"strings"
+)
+
+func SplitHostPort(addr string) (string, string, error) {
+	if net.ParseIP(addr) != nil {
+		return addr, "", nil
+	}
+	host, port, err := net.SplitHostPort(addr)
+	if err == nil {
+		return host, port, nil
+	}
+	if strings.Contains(err.Error(), "missing port") {
+		return addr, "", nil
+	}
+	return addr, "", err
+}
-- 
2.44.0