feat: add frontendURL env

Concurrent map read and map write crashing the whole server

Bug fixes:
- Fix data race in registry.Update
Reviewed-on: #156
Co-authored-by: Bagas <bagas@fossy.my.id>
Co-committed-by: Bagas <bagas@fossy.my.id>

README.md

@@ -36,6 +36,7 @@ The following environment variables can be configured in the `.env` file:
 | Variable            | Description                                                                 | Default                 | Required            |
 |---------------------|-----------------------------------------------------------------------------|-------------------------|---------------------|
 | `DOMAIN`            | Domain name for subdomain routing                                           | `localhost`             | No                  |
+| `FRONTEND_URL`      | URL for the frontend dashboard/landing page                                 | `https://<DOMAIN>`      | No                  |
 | `PORT`              | SSH server port                                                             | `2200`                  | No                  |
 | `HTTP_PORT`         | HTTP server port                                                            | `8080`                  | No                  |
 | `HTTPS_PORT`        | HTTPS server port                                                           | `8443`                  | No                  |

internal/bootstrap/bootstrap_test.go

@@ -80,6 +80,7 @@ type MockConfig struct {
 }
 
 func (m *MockConfig) Domain() string            { return m.Called().String(0) }
+func (m *MockConfig) FrontendURL() string       { return m.Called().String(0) }
 func (m *MockConfig) SSHPort() string           { return m.Called().String(0) }
 func (m *MockConfig) HTTPPort() string          { return m.Called().String(0) }
 func (m *MockConfig) HTTPSPort() string         { return m.Called().String(0) }

internal/config/config.go

@@ -6,6 +6,7 @@ import (
 
 type Config interface {
 	Domain() string
+	FrontendURL() string
 	SSHPort() string
 
 	HTTPPort() string
@@ -50,6 +51,7 @@ func MustLoad() (Config, error) {
 }
 
 func (c *config) Domain() string            { return c.domain }
+func (c *config) FrontendURL() string       { return c.frontendURL }
 func (c *config) SSHPort() string           { return c.sshPort }
 func (c *config) HTTPPort() string          { return c.httpPort }
 func (c *config) HTTPSPort() string         { return c.httpsPort }

internal/config/loader.go

@@ -12,8 +12,9 @@ import (
 )
 
 type config struct {
-	domain  string
+	domain      string
-	sshPort string
+	frontendURL string
+	sshPort     string
 
 	httpPort  string
 	httpsPort string
@@ -49,6 +50,7 @@ func parse() (*config, error) {
 	}
 
 	domain := getenv("DOMAIN", "localhost")
+	frontendURL := getenv("FRONTEND_URL", "https://"+domain)
 	sshPort := getenv("PORT", "2200")
 
 	httpPort := getenv("HTTP_PORT", "8080")
@@ -89,6 +91,7 @@ func parse() (*config, error) {
 
 	return &config{
 		domain:            domain,
+		frontendURL:       frontendURL,
 		sshPort:           sshPort,
 		httpPort:          httpPort,
 		httpsPort:         httpsPort,

internal/grpc/client/client_test.go

@@ -753,6 +753,7 @@ type MockConfig struct {
 }
 
 func (m *MockConfig) Domain() string            { return m.Called().String(0) }
+func (m *MockConfig) FrontendURL() string       { return m.Called().String(0) }
 func (m *MockConfig) SSHPort() string           { return m.Called().String(0) }
 func (m *MockConfig) HTTPPort() string          { return m.Called().String(0) }
 func (m *MockConfig) HTTPSPort() string         { return m.Called().String(0) }

internal/registry/registry.go

@@ -94,13 +94,13 @@ func (r *registry) Update(user string, oldKey, newKey Key) error {
 		return ErrInvalidSlug
 	}
 
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
 	if _, exists := r.slugIndex[newKey]; exists && newKey != oldKey {
 		return ErrSlugInUse
 	}
 
-	r.mu.Lock()
-	defer r.mu.Unlock()
-
 	client, ok := r.byUser[user][oldKey]
 	if !ok {
 		return ErrSessionNotFound

internal/server/server_test.go

@@ -33,6 +33,7 @@ type MockConfig struct {
 }
 
 func (m *MockConfig) Domain() string            { return m.Called().String(0) }
+func (m *MockConfig) FrontendURL() string       { return m.Called().String(0) }
 func (m *MockConfig) SSHPort() string           { return m.Called().String(0) }
 func (m *MockConfig) HTTPPort() string          { return m.Called().String(0) }
 func (m *MockConfig) HTTPSPort() string         { return m.Called().String(0) }

internal/session/forwarder/forwarder_test.go

@@ -24,6 +24,7 @@ type mockConfig struct {
 }
 
 func (m *mockConfig) Domain() string            { return m.Called().String(0) }
+func (m *mockConfig) FrontendURL() string       { return m.Called().String(0) }
 func (m *mockConfig) SSHPort() string           { return m.Called().String(0) }
 func (m *mockConfig) HTTPPort() string          { return m.Called().String(0) }
 func (m *mockConfig) HTTPSPort() string         { return m.Called().String(0) }

internal/session/interaction/interaction_test.go

@@ -32,6 +32,7 @@ type MockConfig struct {
 }
 
 func (m *MockConfig) Domain() string            { return m.Called().String(0) }
+func (m *MockConfig) FrontendURL() string       { return m.Called().String(0) }
 func (m *MockConfig) SSHPort() string           { return m.Called().String(0) }
 func (m *MockConfig) HTTPPort() string          { return m.Called().String(0) }
 func (m *MockConfig) HTTPSPort() string         { return m.Called().String(0) }

internal/session/session.go

@@ -363,19 +363,28 @@ func (s *session) HandleTCPForward(req *ssh.Request, addr string, portToBind uin
 		}
 	}
 
+	releasePort := func() {
+		if err := s.lifecycle.PortRegistry().SetStatus(portToBind, false); err != nil {
+			log.Printf("failed to release port %d: %v", portToBind, err)
+		}
+	}
+
 	tcpServer := transport.NewTCPServer(portToBind, s.forwarder)
 	listener, err := tcpServer.Listen()
 	if err != nil {
+		releasePort()
 		return s.denyForwardingRequest(req, nil, listener, fmt.Sprintf("Port %d is already in use or restricted", portToBind))
 	}
 
 	key := types.SessionKey{Id: fmt.Sprintf("%d", portToBind), Type: types.TunnelTypeTCP}
 	if !s.registry.Register(key, s) {
+		releasePort()
 		return s.denyForwardingRequest(req, nil, listener, fmt.Sprintf("Failed to register TunnelTypeTCP client with id: %s", key.Id))
 	}
 
 	err = s.finalizeForwarding(req, portToBind, listener, types.TunnelTypeTCP, key.Id)
 	if err != nil {
+		releasePort()
 		return s.denyForwardingRequest(req, &key, listener, fmt.Sprintf("Failed to finalize forwarding: %s", err))
 	}
 

internal/session/session_test.go

@@ -38,8 +38,9 @@ type mockConfig struct {
 	config.Config
 }
 
-func (m *mockConfig) Domain() string  { return m.Called().String(0) }
+func (m *mockConfig) Domain() string            { return m.Called().String(0) }
-func (m *mockConfig) SSHPort() string { return m.Called().String(0) }
+func (m *mockConfig) FrontendURL() string       { return m.Called().String(0) }
+func (m *mockConfig) SSHPort() string           { return m.Called().String(0) }
 func (m *mockConfig) Mode() types.ServerMode {
 	args := m.Called()
 	if args.Get(0) == nil {
@@ -705,6 +706,7 @@ func TestForwardingFailures(t *testing.T) {
 		s, mRegistry, mPort, _, _, sReqs, cConn, cleanup := setup(t)
 		defer cleanup()
 		mPort.On("Claim", mock.Anything).Return(true)
+		mPort.On("SetStatus", uint16(1234), false).Return(nil)
 		mRegistry.On("Register", mock.Anything, mock.Anything).Return(false)
 
 		payload := make([]byte, 4+9+4)
@@ -723,7 +725,7 @@ func TestForwardingFailures(t *testing.T) {
 	})
 
 	t.Run("Finalize Forwarding Failure", func(t *testing.T) {
-		s, mRegistry, _, mRandom, _, sReqs, cConn, cleanup := setup(t)
+		s, mRegistry, _, mRandom, sConn, sReqs, cConn, cleanup := setup(t)
 		defer cleanup()
 		mRandom.On("String", 20).Return("test-slug", nil)
 		mRegistry.On("Register", mock.Anything, mock.Anything).Return(true)
@@ -742,7 +744,7 @@ func TestForwardingFailures(t *testing.T) {
 		err := cConn.Close()
 		assert.NoError(t, err)
 
-		time.Sleep(50 * time.Millisecond)
+		_ = sConn.Wait()
 
 		err = s.HandleTCPIPForward(req)
 		assert.Error(t, err)
@@ -764,6 +766,7 @@ func TestForwardingFailures(t *testing.T) {
 		}(l)
 		_, portStr, _ := net.SplitHostPort(l.Addr().String())
 		port, _ := strconv.Atoi(portStr)
+		mPort.On("SetStatus", uint16(port), false).Return(nil)
 
 		payload := make([]byte, 4+9+4)
 		binary.BigEndian.PutUint32(payload[0:4], 9)
@@ -1247,6 +1250,7 @@ func TestHandleTCPForward_Failures(t *testing.T) {
 			assert.NoError(t, err)
 		}(l)
 		port := uint16(l.Addr().(*net.TCPAddr).Port)
+		mPort.On("SetStatus", port, false).Return(nil)
 
 		err = s.HandleTCPForward(getReq(t, cConn, sReqs), "localhost", port, false)
 		if err == nil {
@@ -1254,12 +1258,14 @@ func TestHandleTCPForward_Failures(t *testing.T) {
 		} else if !strings.Contains(err.Error(), "already in use") {
 			t.Errorf("expected error to contain %q, got %q", "already in use", err.Error())
 		}
+		mPort.AssertExpectations(t)
 	})
 
 	t.Run("Registry Register fail", func(t *testing.T) {
 		s, mRegistry, mPort, _, sReqs, cConn, cleanup := setup(t)
 		defer cleanup()
 		mPort.On("Claim", mock.Anything).Return(true)
+		mPort.On("SetStatus", mock.AnythingOfType("uint16"), false).Return(nil)
 		mRegistry.On("Register", mock.Anything, mock.Anything).Return(false)
 		err := s.HandleTCPForward(getReq(t, cConn, sReqs), "localhost", 0, false)
 		if err == nil {
@@ -1267,12 +1273,14 @@ func TestHandleTCPForward_Failures(t *testing.T) {
 		} else if !strings.Contains(err.Error(), "Failed to register") {
 			t.Errorf("expected error to contain %q, got %q", "Failed to register", err.Error())
 		}
+		mPort.AssertExpectations(t)
 	})
 
 	t.Run("Finalize fail (Reply fail)", func(t *testing.T) {
 		s, mRegistry, mPort, sConn, sReqs, cConn, cleanup := setup(t)
 		defer cleanup()
 		mPort.On("Claim", mock.Anything).Return(true)
+		mPort.On("SetStatus", mock.AnythingOfType("uint16"), false).Return(nil)
 		mRegistry.On("Register", mock.Anything, mock.Anything).Return(true)
 		req := getReq(t, cConn, sReqs)
 		err := cConn.Close()
@@ -1285,6 +1293,7 @@ func TestHandleTCPForward_Failures(t *testing.T) {
 		} else if !strings.Contains(err.Error(), "Failed to finalize forwarding") {
 			t.Errorf("expected error to contain %q, got %q", "Failed to finalize forwarding", err.Error())
 		}
+		mPort.AssertExpectations(t)
 	})
 }
 

internal/transport/httphandler.go

@@ -116,7 +116,7 @@ func (hh *httpHandler) Handler(conn net.Conn, isTLS bool) {
 		Type: types.TunnelTypeHTTP,
 	})
 	if err != nil {
-		_ = hh.redirect(conn, http.StatusMovedPermanently, fmt.Sprintf("https://tunnl.live/tunnel-not-found?slug=%s\r\n", slug))
+		_ = hh.redirect(conn, http.StatusMovedPermanently, fmt.Sprintf("%s/tunnel-not-found?slug=%s\r\n", hh.config.FrontendURL(), slug))
 		return
 	}
 

internal/transport/httphandler_test.go

@@ -223,6 +223,7 @@ func TestNewHTTPHandler(t *testing.T) {
 	msr := new(MockSessionRegistry)
 	mockConfig := &MockConfig{}
 	mockConfig.On("Domain").Return("domain")
+	mockConfig.On("FrontendURL").Return("https://domain")
 	mockConfig.On("TLSRedirect").Return(false)
 	hh := newHTTPHandler(mockConfig, msr)
 	assert.NotNil(t, hh)
@@ -290,7 +291,7 @@ func TestHandler(t *testing.T) {
 			isTLS:       true,
 			redirectTLS: false,
 			request:     []byte("GET / HTTP/1.1\r\nHost: test.domain\r\n\r\n"),
-			expected:    []byte("HTTP/1.1 301 Moved Permanently\r\nLocation: https://tunnl.live/tunnel-not-found?slug=test\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"),
+			expected:    []byte("HTTP/1.1 301 Moved Permanently\r\nLocation: https://example.com/tunnel-not-found?slug=test\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"),
 			setupMocks: func(msr *MockSessionRegistry) {
 				msr.On("Get", types.SessionKey{
 					Id:   "test",
@@ -616,6 +617,7 @@ func TestHandler(t *testing.T) {
 			mockConfig := &MockConfig{}
 			port := "0"
 			mockConfig.On("Domain").Return("example.com")
+			mockConfig.On("FrontendURL").Return("https://example.com")
 			mockConfig.On("HTTPPort").Return(port)
 			mockConfig.On("HeaderSize").Return(4096)
 			mockConfig.On("TLSRedirect").Return(true)
@@ -726,6 +728,7 @@ func TestHandlerForwardsPostBody(t *testing.T) {
 	mockSessionRegistry := new(MockSessionRegistry)
 	mockConfig := &MockConfig{}
 	mockConfig.On("Domain").Return("example.com")
+	mockConfig.On("FrontendURL").Return("https://example.com")
 	mockConfig.On("HTTPPort").Return("0")
 	mockConfig.On("HeaderSize").Return(4096)
 	mockConfig.On("TLSRedirect").Return(true)

internal/transport/tls_test.go

@@ -25,6 +25,7 @@ type MockConfig struct {
 }
 
 func (m *MockConfig) Domain() string            { return m.Called().String(0) }
+func (m *MockConfig) FrontendURL() string       { return m.Called().String(0) }
 func (m *MockConfig) SSHPort() string           { return m.Called().String(0) }
 func (m *MockConfig) HTTPPort() string          { return m.Called().String(0) }
 func (m *MockConfig) HTTPSPort() string         { return m.Called().String(0) }