serve http response for subdomain not found and client not online

Author: amalshaji

internal/client/config/config.go

@@ -58,9 +58,6 @@ func (c *Config) SetDefaults() {
 	for i := range c.Tunnels {
 		c.Tunnels[i].setDefaults()
 	}
-
-	c.Secure = true
-	c.Debug = true
 }
 
 type ClientConfig struct {

internal/client/ssh/ssh.go

@@ -9,6 +9,7 @@ import (
 	"net"
 	"os"
 	"os/signal"
+	"strings"
 	"syscall"
 	"time"
 
@@ -72,7 +73,6 @@ func (s *SshClient) startListenerForClient() error {
 	}
 	defer sshClient.Close()
 
-	// remoteEndpoint := "localhost:0"            // Remote address to listen on
 	localEndpoint := s.config.Tunnel.GetAddr() // Local address to forward to
 
 	randomPorts := utils.GenerateRandomHttpPorts()[:1]
@@ -98,13 +98,15 @@ func (s *SshClient) startListenerForClient() error {
 		remoteConn, err := s.listener.Accept()
 		if err != nil {
 			s.log.Error("failed to accept connection", "error", err)
-			continue
+			break
 		}
 
 		// Connect to the local endpoint
 		localConn, err := net.Dial("tcp", localEndpoint)
 		if err != nil {
-			s.log.Error("failed to connect to local endpoint", "error", err)
+			// serve local html if the local server is not available
+			// change this to beautiful template
+			remoteConn.Write([]byte(strings.TrimSpace(utils.LocalServerNotOnline)))
 			remoteConn.Close()
 			continue
 		}
@@ -113,6 +115,8 @@ func (s *SshClient) startListenerForClient() error {
 		go tunnel(remoteConn, localConn)
 		go tunnel(localConn, remoteConn)
 	}
+
+	return nil
 }
 
 func tunnel(src, dst net.Conn) {

internal/server/proxy/proxy.go

@@ -10,6 +10,7 @@ import (
 	"net/url"
 	"os"
 	"os/signal"
+	"strings"
 	"sync"
 	"syscall"
 	"time"
@@ -26,13 +27,19 @@ type Proxy struct {
 	server *http.Server
 }
 
+func (p *Proxy) GetServerAddr() string {
+	return ":" + fmt.Sprint(p.config.Proxy.Port)
+}
+
 func New(config *config.Config) *Proxy {
-	return &Proxy{
+	p := &Proxy{
 		log:    utils.GetLogger(),
 		config: config,
 		routes: make(map[string]string),
-		server: &http.Server{Addr: ":" + fmt.Sprint(config.Proxy.Port)},
+		server: nil,
 	}
+	p.server = &http.Server{Addr: p.GetServerAddr()}
+	return p
 }
 
 func (p *Proxy) GetRoute(src string) (string, error) {
@@ -77,7 +84,8 @@ func (p *Proxy) handleRequest(w http.ResponseWriter, r *http.Request) {
 	subdomain := p.config.ExtractSubdomain(r.Host)
 	target, err := p.GetRoute(subdomain)
 	if err != nil {
-		http.Error(w, err.Error(), http.StatusNotFound)
+		w.WriteHeader(http.StatusNotFound)
+		w.Write([]byte(strings.TrimSpace(utils.UnregisteredSubdomain)))
 		return
 	}
 
@@ -90,7 +98,7 @@ func (p *Proxy) handleRequest(w http.ResponseWriter, r *http.Request) {
 }
 
 func (p *Proxy) Start() {
-	p.log.Info("starting proxy server", "host", p.config.Proxy.Host, "port", p.config.Proxy.Port)
+	p.log.Info("starting proxy server", "port", p.GetServerAddr())
 
 	done := make(chan os.Signal, 1)
 	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)

internal/server/ssh/sshd.go

@@ -32,6 +32,10 @@ func New(config *config.SshConfig, proxy *proxy.Proxy) *SshServer {
 	}
 }
 
+func (s *SshServer) GetServerAddr() string {
+	return ":" + fmt.Sprint(s.config.Port)
+}
+
 func (s *SshServer) getSshPublicKey() ssh.PublicKey {
 	publicKey, err := os.ReadFile(s.config.KeysDir + "/id_rsa.pub")
 	if err != nil {
@@ -41,13 +45,11 @@ func (s *SshServer) getSshPublicKey() ssh.PublicKey {
 	return key
 }
 
-type Status struct {
-	user   string
-	status string
-}
-
 func parseSshUsername(user string) (string, string) {
 	output := strings.Split(user, ":")
+	if len(output) != 2 {
+		log.Fatal("invalid username format")
+	}
 	return output[0], output[1]
 }
 
@@ -56,16 +58,8 @@ func (s *SshServer) Start() {
 
 	keyFromDisk := s.getSshPublicKey()
 
-	status := make(chan Status)
-
-	go func() {
-		for st := range status {
-			s.log.Info("connection status", "user", st.user, "status", st.status)
-		}
-	}()
-
 	server := ssh.Server{
-		Addr: ":" + fmt.Sprint(s.config.Port),
+		Addr: s.GetServerAddr(),
 		Handler: ssh.Handler(func(sh ssh.Session) {
 			select {}
 		}),
@@ -85,6 +79,10 @@ func (s *SshServer) Start() {
 				<-ctx.Done()
 				// log user disconnection
 				log.Printf("connection closed for %s", user)
+				err := s.proxy.RemoveRoute(subdomain)
+				if err != nil {
+					s.log.Error("failed to remove route", "error", err)
+				}
 			}()
 			return true
 		}),
@@ -101,7 +99,7 @@ func (s *SshServer) Start() {
 	done := make(chan os.Signal, 1)
 	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
 
-	s.log.Info("starting SSH server", "host", s.config.Host, "port", s.config.Port)
+	s.log.Info("starting SSH server", "port", s.GetServerAddr())
 
 	go func() {
 		if err := server.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {

internal/utils/http.go

@@ -0,0 +1,27 @@
+package utils
+
+const LocalServerNotOnline = `
+HTTP/1.1 400 Bad Request
+Content-Type: text/html; charset=utf-8
+
+<!DOCTYPE html>
+<html>
+<head>
+    <title>Localport</title>
+</head>
+<body>
+    <h1>Your local server is not online</h1>
+	<h2>Try again after starting the server</h2>
+</body>
+</html>`
+
+const UnregisteredSubdomain = `
+<!DOCTYPE html>
+<html>
+<head>
+    <title>Localport</title>
+</head>
+<body>
+    <h1>There is no tunnel running on this endpoint</h1>
+</body>
+</html>`