Add CNAME Support

azimjohn · Jan 2, 2024 · bae161b · bae161b
1 parent f5b106f
commit bae161b
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 9 deletions.
diff --git a/cli/jprqc.go b/cli/jprqc.go
@@ -17,6 +17,7 @@ type jprqClient struct {
 	config       Config
 	protocol     string
 	subdomain    string
+	cname        string
 	localServer  string
 	remoteServer string
 	publicServer string
@@ -34,6 +35,7 @@ func (j *jprqClient) Start(port int, debug bool) {
 		Data: &events.TunnelRequested{
 			Protocol:   j.protocol,
 			Subdomain:  j.subdomain,
+			CanonName:  j.cname,
 			AuthToken:  j.config.Local.AuthToken,
 			CliVersion: version,
 		},
@@ -42,24 +44,24 @@ func (j *jprqClient) Start(port int, debug bool) {
 		log.Fatalf("failed to send request: %s\n", err)
 	}
 
-	var tunnel events.Event[events.TunnelOpened]
-	if err := tunnel.Read(eventCon); err != nil {
+	var t events.Event[events.TunnelOpened]
+	if err := t.Read(eventCon); err != nil {
 		log.Fatalf("failed to receive tunnel info: %s\n", err)
 	}
-	if tunnel.Data.ErrorMessage != "" {
-		log.Fatalf(tunnel.Data.ErrorMessage)
+	if t.Data.ErrorMessage != "" {
+		log.Fatalf(t.Data.ErrorMessage)
 	}
 
 	j.localServer = fmt.Sprintf("127.0.0.1:%d", port)
-	j.remoteServer = fmt.Sprintf("jprq.%s:%d", j.config.Remote.Domain, tunnel.Data.PrivateServer)
-	j.publicServer = fmt.Sprintf("%s:%d", tunnel.Data.Hostname, tunnel.Data.PublicServer)
+	j.remoteServer = fmt.Sprintf("jprq.%s:%d", j.config.Remote.Domain, t.Data.PrivateServer)
+	j.publicServer = fmt.Sprintf("%s:%d", t.Data.Hostname, t.Data.PublicServer)
 
 	fmt.Printf("Status: \t Online \n")
 	fmt.Printf("Protocol: \t %s \n", strings.ToUpper(j.protocol))
 	fmt.Printf("Forwarded: \t %s -> %s \n", strings.TrimSuffix(j.publicServer, ":80"), j.localServer)
 
 	if j.protocol == "http" {
-		j.publicServer = fmt.Sprintf("https://%s", tunnel.Data.Hostname)
+		j.publicServer = fmt.Sprintf("https://%s", t.Data.Hostname)
 	}
 
 	if j.protocol == "http" && debug {

diff --git a/cli/main.go b/cli/main.go
@@ -16,6 +16,7 @@ var version = "2.2"
 
 type Flags struct {
 	debug     bool
+	cname     string
 	subdomain string
 }
 
@@ -89,6 +90,7 @@ func main() {
 		config:    conf,
 		protocol:  protocol,
 		subdomain: flags.subdomain,
+		cname:     flags.cname,
 	}
 
 	go client.Start(port, flags.debug)
@@ -102,10 +104,12 @@ func parseFlags(args []string) Flags {
 	var flags Flags
 	for i, arg := range args {
 		switch arg {
-		case "-debug", "--debug":
+		case "-d", "-debug", "--debug":
 			flags.debug = true
-		case "-s", "--subdomain":
+		case "-s", "-subdomain", "--subdomain":
 			flags.subdomain = args[i+1]
+		case "-c", "-cname", "--cname":
+			flags.cname = args[i+1]
 		}
 	}
 	return flags

diff --git a/server/events/events.go b/server/events/events.go
@@ -26,6 +26,7 @@ type Event[Type EventType] struct {
 type TunnelRequested struct {
 	Protocol   string
 	Subdomain  string
+	CanonName  string
 	AuthToken  string
 	CliVersion string
 }

diff --git a/server/jprq.go b/server/jprq.go
@@ -28,6 +28,7 @@ type Jprq struct {
 	allowedUsers    map[string]string
 	allowedLastMod  time.Time
 	authenticator   github.Authenticator
+	cnameMap        map[string]string
 	tcpTunnels      map[uint16]*tunnel.TCPTunnel
 	httpTunnels     map[string]*tunnel.HTTPTunnel
 	userTunnels     map[string]map[string]tunnel.Tunnel
@@ -37,6 +38,7 @@ func (j *Jprq) Init(conf config.Config, oauth github.Authenticator) error {
 	j.config = conf
 	j.authenticator = oauth
 	j.allowedUsers = make(map[string]string)
+	j.cnameMap = make(map[string]string)
 	j.tcpTunnels = make(map[uint16]*tunnel.TCPTunnel)
 	j.httpTunnels = make(map[string]*tunnel.HTTPTunnel)
 	j.userTunnels = make(map[string]map[string]tunnel.Tunnel)
@@ -87,6 +89,9 @@ func (j *Jprq) servePublicConn(conn net.Conn) error {
 		writeResponse(conn, 400, "Bad Request", "Bad Request")
 		return nil
 	}
+	if tunnelHost, ok := j.cnameMap[host]; ok && tunnelHost != "" {
+		host = tunnelHost
+	}
 	host = strings.ToLower(host)
 	t, found := j.httpTunnels[host]
 	if !found {
@@ -129,6 +134,10 @@ func (j *Jprq) serveEventConn(conn net.Conn) error {
 	if _, ok := j.httpTunnels[hostname]; ok {
 		return events.WriteError(conn, "subdomain is busy: %s, try another one", request.Subdomain)
 	}
+	cname := request.CanonName
+	if _, ok := j.cnameMap[cname]; ok && cname != "" {
+		return events.WriteError(conn, "cname is busy: %s, try another one", request.CanonName)
+	}
 
 	var t tunnel.Tunnel
 	var maxConsLimit = j.config.MaxConsPerTunnel
@@ -139,7 +148,9 @@ func (j *Jprq) serveEventConn(conn net.Conn) error {
 		if err != nil {
 			return events.WriteError(conn, "failed to create http tunnel", err.Error())
 		}
+		j.cnameMap[cname] = hostname
 		j.httpTunnels[hostname] = tn
+		defer delete(j.cnameMap, cname)
 		defer delete(j.httpTunnels, hostname)
 		t = tn
 	case events.TCP: