grokway/cmd/server/main.go

package main

import (
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"flag"

	"github.com/gliderlabs/ssh"
	gossh "golang.org/x/crypto/ssh"
)

const charset = "abcdefghijklmnopqrstuvwxyz0123456789"

func generateSlug(length int) string {
	b := make([]byte, length)
	for i := range b {
		b[i] = charset[rand.Intn(len(charset))]
	}
	return string(b)
}

// TunnelManager manages the active tunnels
type TunnelManager struct {
	tunnels map[string]*Tunnel
	mu      sync.RWMutex
}

type Tunnel struct {
	ID        string
	LocalPort uint32
	Conn      *gossh.ServerConn
}

var manager = &TunnelManager{
	tunnels: make(map[string]*Tunnel),
}

func (tm *TunnelManager) Register(id string, t *Tunnel) {
	tm.mu.Lock()
	defer tm.mu.Unlock()
	tm.tunnels[id] = t
}

func (tm *TunnelManager) Unregister(id string) {
	tm.mu.Lock()
	defer tm.mu.Unlock()
	delete(tm.tunnels, id)
}

func (tm *TunnelManager) Get(id string) (*Tunnel, bool) {
	tm.mu.RLock()
	defer tm.mu.RUnlock()
	t, ok := tm.tunnels[id]
	return t, ok
}

func (tm *TunnelManager) FindSlugByConn(conn *gossh.ServerConn) (string, bool) {
	tm.mu.RLock()
	defer tm.mu.RUnlock()
	for slug, t := range tm.tunnels {
		if t.Conn == conn {
			return slug, true
		}
	}
	return "", false
}

// HTTP Proxy to forward requests to the SSH tunnel
func startHttpProxy(port string) {
	log.Printf("Starting HTTP Proxy on %s", port)
	// Serve static files for 404 page

	http.Handle("/assets/", http.StripPrefix("/assets/", http.FileServer(http.Dir("./cmd/server/static/assets"))))

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Extract subdomain to identify the tunnel
		// Format expects: slug.grokway.com or slug.localhost
		host := r.Host
		parts := strings.Split(host, ".")
		// If accessing IP directly or localhost without subdomain
		if len(parts) < 2 || host == "localhost:8080" {
			http.Error(w, "Grokway Server Running", http.StatusOK)
			return
		}
		slug := parts[0]

		tunnel, ok := manager.Get(slug)
		if !ok {
			// Serve 404 page
			http.ServeFile(w, r, "./cmd/server/static/404.html")
			return
		}

		// In a real implementation, we would dial the tunnel
		// using ssh.Conn.OpenChannel("direct-tcpip", ...)
		// But gliderlabs/ssh handles Reverse Port Forwarding differently.
		// It usually binds a listener on the server side.

		// Wait! The standard "Remote Port Forwarding" (-R) means the SSH Server
		// listens on a port ON THE SERVER, and forwards traffic to the client.

		// If GrokwayV1 used -R 80:localhost:8080, then the server was listening on 80.
		// Here we want dynamic subdomains.
		// So we likely want to use a Custom Request or just standard forwarding
		// but mapped to our HTTP proxy.

		// SIMPLE APPROACH (Like Ngrok):
		// 1. Client connects via SSH.
		// 2. Client requests "tcpip-forward" (remote forwarding).
		// 3. Server acknowledges.
		// 4. When HTTP request comes to `slug.grokway.com`:
		//    a. Server accepts request.
		//    b. Server opens a "forwarded-tcpip" channel TO the client.
		//    c. Server proxies the HTTP request body over this channel.

		// For this MVP, let's assume we implement the "forwarded-tcpip" logic.
		// However, gliderlabs/ssh simplifies this if we use its Forwarding support.
		// But typically `gliderlabs/ssh` is for allowing the server to be a jump host.
		// We want to be an HTTP Gateway.

// 1. Hijack the connection to handle bidirectional traffic (WebSockets)
		hijacker, ok := w.(http.Hijacker)
		if !ok {
			http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
			return
		}
		clientConn, bufrw, err := hijacker.Hijack()
		if err != nil {
			http.Error(w, err.Error(), http.StatusServiceUnavailable)
			return
		}
		defer clientConn.Close()

		// 2. Open channel to client
		// "forwarded-tcpip" arguments:
		destHost := "0.0.0.0"
		destPort := tunnel.LocalPort
		srcHost := "127.0.0.1"
		srcPort := uint32(12345)

		payload := make([]byte, 0)
		payload = append(payload, []byte{0, 0, 0, uint8(len(destHost))}...)
		payload = append(payload, []byte(destHost)...)
		payload = append(payload, byte(destPort>>24), byte(destPort>>16), byte(destPort>>8), byte(destPort))
		payload = append(payload, []byte{0, 0, 0, uint8(len(srcHost))}...)
		payload = append(payload, []byte(srcHost)...)
		payload = append(payload, byte(srcPort>>24), byte(srcPort>>16), byte(srcPort>>8), byte(srcPort))

		ch, reqs, err := tunnel.Conn.OpenChannel("forwarded-tcpip", payload)
		if err != nil {
			log.Printf("Failed to open forwarded-tcpip channel: %s", err)
			return
		}
		defer ch.Close()
		go gossh.DiscardRequests(reqs)

		var wg sync.WaitGroup
		wg.Add(2)

		// 3. Browser -> Backend (Write request + Copy raw stream)
		go func() {
			defer wg.Done()
			// Write the initial request (Method, Path, Headers)
			// This sets up the handshake or request.
			// Note: We use r.Write to reconstruct the request line and headers.
			// For WebSockets, the Body is empty, so this writes headers and returns.
			// For POSTs, it writes headers and tries to copy Body.
			if err := r.Write(ch); err != nil {
				log.Printf("Error writing request to backend: %v", err)
				return
			}
			// Important: Continue copying any subsequent data (like WebSocket frames)
			// from the hijacked buffer/connection to the channel.
			io.Copy(ch, bufrw)
			// e.g. when browser closes or stops sending, we are done here.
		}()

		// 4. Backend -> Browser (Copy raw stream)
		go func() {
			defer wg.Done()
			io.Copy(clientConn, ch)
			// When backend closes connection, close browser connection
			clientConn.Close()
		}()

		wg.Wait()

	})

	log.Fatal(http.ListenAndServe(port, nil))
}

func main() {
	rand.Seed(time.Now().UnixNano())

	// 0. Parse Flags
	sshAddr := flag.String("ssh", ":2222", "Address for SSH server to listen on")
	httpAddr := flag.String("http", ":8080", "Address for HTTP proxy to listen on")
	flag.Parse()

	// 1. Auth Token Setup
	authToken := os.Getenv("GROKWAY_TOKEN")
	if authToken == "" {
		authToken = generateSlug(16) // Longer token
		fmt.Printf("\nGenerated Admin Token: %s\n\n", authToken)
		fmt.Println("Use this token to configure your client:")
		fmt.Printf("  grokway client configure --token %s --server <GROKWAY_SERVER_ADDRESS>\n\n", authToken)
	} else {
		fmt.Println("Using configured GROKWAY_TOKEN")
	}

	// 2. SSH Server
	sshServer := &ssh.Server{
		Addr: *sshAddr,
		PasswordHandler: func(ctx ssh.Context, password string) bool {
			return password == authToken
		},
		Handler: func(s ssh.Session) {
			// This handler handles the "Shell" or "Exec" requests.
			// The port forwarding happens via "tcpip-forward" request global handler.

			// Show TUI-ish banner to the SSH client
			io.WriteString(s, "Welcome to GrokwayV2!\n")
			io.WriteString(s, fmt.Sprintf("You are connected as %s\n", s.User()))

			// Keep session open
			select {}
		},
	}

	// Register global request handler for "tcpip-forward"
	// This is where clients say "Please listen on port X"
	sshServer.RequestHandlers = map[string]ssh.RequestHandler{
		"tcpip-forward": func(ctx ssh.Context, srv *ssh.Server, req *gossh.Request) (bool, []byte) {
			// Parse payload
			// string address to bind (usually empty or 0.0.0.0 or 127.0.0.1)
			// uint32 port number to bind

			// For Grokway, we ignore the requested port and assign a random subdomain/slug
			// Or we use the requested port if valid?
			// Let's assume we ignore it and generate a slug,
			// OR we use the port as the "slug" if it's special.

			// But wait, the client needs to know what we assigned!
			// Standard SSH response to tcpip-forward contains the BOUND PORT.
			// If we return a port, the client knows.
			// If we want to return a URL, standard SSH doesn't have a field for that.
			// But we can write it to the Session (stdout).

			// Let's accept any request and map it to a random slug.
			slug := generateSlug(8)

			log.Printf("Client requested forwarding. Assigning slug: %s", slug)

			conn := ctx.Value(ssh.ContextKeyConn).(*gossh.ServerConn)

			manager.Register(slug, &Tunnel{
				ID:        slug,
				LocalPort: 80, // We assume client forwards to port 80 locally? No,
				// the client says "forward FROM remote port X".
				// The server opens channel TO client when traffic hits X.
				Conn: conn,
			})

			// Monitor connection for closure to cleanup
			go func() {
				conn.Wait()
				log.Printf("Connection closed for slug %s, unregistering...", slug)
				manager.Unregister(slug)
			}()

			// Reply success
			// Payload: port that was bound (uint32)
			// We can return a successful port, e.g. 0 or 80.
			portBytes := []byte{0, 0, 0, 80}
			return true, portBytes
		},
		"cancel-tcpip-forward": func(ctx ssh.Context, srv *ssh.Server, req *gossh.Request) (bool, []byte) {
			return true, nil
		},
		"grokway-whoami": func(ctx ssh.Context, srv *ssh.Server, req *gossh.Request) (bool, []byte) {
			conn := ctx.Value(ssh.ContextKeyConn).(*gossh.ServerConn)
			slug, ok := manager.FindSlugByConn(conn)
			if !ok {
				// Not found? Maybe not forwarded yet?
				log.Printf("Whoami requested but no tunnel found for conn %v", conn.RemoteAddr())
				return false, nil
			}
			return true, []byte(slug)
		},
	}

	// Load Host Key (optional for dev, helpful for no warnings)
	// sshServer.SetOption(ssh.HostKeyFile("host.key"))

	fmt.Printf("Starting GrokwayV2 Server on %s\n", *sshAddr)
	fmt.Printf("Starting HTTP Proxy on %s\n", *httpAddr)

	go startHttpProxy(*httpAddr)

	log.Fatal(sshServer.ListenAndServe())
}