+ New

fierj

Public

Tiny personal git forge

← fierj / ssh.go
package main

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
	"fmt"
	"io"
	"log/slog"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"
)

type SSHServer struct {
	reposDir string
	config   *ssh.ServerConfig
}

func NewSSHServer(reposDir, hostKeyPath string) (*SSHServer, error) {
	signer, err := loadOrGenerateHostKey(hostKeyPath)
	if err != nil {
		return nil, fmt.Errorf("host key: %w", err)
	}

	config := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			slog.Info("ssh connection accepted", "user", conn.User(), "remote", conn.RemoteAddr())
			return &ssh.Permissions{
				Extensions: map[string]string{"pubkey": string(ssh.MarshalAuthorizedKey(key))},
			}, nil
		},
	}
	config.AddHostKey(signer)

	return &SSHServer{reposDir: reposDir, config: config}, nil
}

func (s *SSHServer) ListenAndServe(addr string) error {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	slog.Info("ssh server listening", "addr", addr)

	for {
		conn, err := listener.Accept()
		if err != nil {
			slog.Error("ssh accept error", "error", err)
			continue
		}
		go s.handleConn(conn)
	}
}

func (s *SSHServer) handleConn(conn net.Conn) {
	defer conn.Close()

	sshConn, chans, reqs, err := ssh.NewServerConn(conn, s.config)
	if err != nil {
		slog.Warn("ssh handshake failed", "error", err, "remote", conn.RemoteAddr())
		return
	}
	defer sshConn.Close()

	go ssh.DiscardRequests(reqs)

	for ch := range chans {
		if ch.ChannelType() != "session" {
			ch.Reject(ssh.UnknownChannelType, "only session channels supported")
			continue
		}
		channel, requests, err := ch.Accept()
		if err != nil {
			continue
		}
		go s.handleSession(channel, requests)
	}
}

func (s *SSHServer) handleSession(channel ssh.Channel, requests <-chan *ssh.Request) {
	defer channel.Close()

	for req := range requests {
		var execMsg struct{ Command string }
		switch req.Type {
		case "exec":
			ssh.Unmarshal(req.Payload, &execMsg)
			req.Reply(true, nil)
			s.runGitCommand(channel, execMsg.Command)
			return
		case "env":
			req.Reply(false, nil)
		default:
			req.Reply(false, nil)
		}
	}
}

func (s *SSHServer) runGitCommand(channel ssh.Channel, gitCmd string) {
	parts := strings.SplitN(strings.TrimSpace(gitCmd), " ", 2)
	if len(parts) != 2 {
		io.WriteString(channel.Stderr(), "fatal: invalid git command\n")
		return
	}

	subCmd := parts[0]
	repoArg := strings.Trim(parts[1], "'\"")
	repoName := strings.TrimSuffix(repoArg, ".git")
	repoName = strings.TrimPrefix(repoName, "/")

	if strings.Contains(repoName, "..") || strings.Contains(repoName, "/") {
		io.WriteString(channel.Stderr(), "fatal: invalid repository path\n")
		return
	}

	absRoot, err := filepath.Abs(s.reposDir)
	if err != nil {
		io.WriteString(channel.Stderr(), "fatal: internal error\n")
		return
	}
	repoPath := filepath.Join(absRoot, repoName+".git")

	if _, err := os.Stat(repoPath); os.IsNotExist(err) {
		io.WriteString(channel.Stderr(), "fatal: repository not found\n")
		return
	}

	switch subCmd {
	case "git-upload-pack", "git-upload-archive", "git-receive-pack":
	default:
		io.WriteString(channel.Stderr(), "fatal: unknown git command: "+subCmd+"\n")
		return
	}

	slog.Info("ssh git", "cmd", subCmd, "repo", repoName)

	cmd := exec.Command(subCmd, ".")
	cmd.Dir = repoPath

	stdin, _ := cmd.StdinPipe()
	stdout, _ := cmd.StdoutPipe()
	cmd.Stderr = channel.Stderr()

	if err := cmd.Start(); err != nil {
		slog.Warn("start failed", "err", err)
		return
	}

	// Stdout -> channel (must drain before we return).
	// Stdin <- channel (will be unblocked by channel.Close in handleSession).
	var wg sync.WaitGroup
	wg.Add(1)
	go func() { defer wg.Done(); io.Copy(channel, stdout) }()
	go func() { io.Copy(stdin, channel); stdin.Close() }()

	if err := cmd.Wait(); err != nil {
		slog.Warn("git failed", "cmd", subCmd, "repo", repoName, "err", err)
	}
	wg.Wait()

	// Send exit status so the client knows the command finished successfully.
	// Without this, git push reports an error even when the push succeeded.
	exitCode := uint32(0)
	if cmd.ProcessState != nil && !cmd.ProcessState.Success() {
		exitCode = 1
	}
	channel.SendRequest("exit-status", false, ssh.Marshal(&struct{ ExitStatus uint32 }{exitCode}))

	// Give the SSH library time to flush before channel.Close in handleSession.
	time.Sleep(50 * time.Millisecond)
}

func loadOrGenerateHostKey(path string) (ssh.Signer, error) {
	if path == "" {
		path = "fierj_host_key"
	}

	if data, err := os.ReadFile(path); err == nil {
		return ssh.ParsePrivateKey(data)
	}

	slog.Info("generating new SSH host key", "path", path)
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, err
	}

	privBytes, err := ssh.MarshalPrivateKey(priv, "")
	if err != nil {
		return nil, err
	}
	if err := os.WriteFile(path, pem.EncodeToMemory(privBytes), 0600); err != nil {
		return nil, err
	}

	signer, err := ssh.NewSignerFromKey(priv)
	if err != nil {
		return nil, err
	}
	return signer, nil
}