Dump UNIX sockets to standard output

unixdump.go

package main

import (
	"fmt"
	"io"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

func main() {
	if len(os.Args) != 2 {
		fmt.Fprintln(os.Stderr, "Usage:", os.Args[0], " <socket-path>")
		os.Exit(1)
	}

	socketPath := os.Args[1]
	origSocketPath := socketPath + ".orig"

	// Setup signal handling for graceful shutdown
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT, syscall.SIGQUIT)

	// Check if the socket exists
	if _, err := os.Stat(socketPath); os.IsNotExist(err) {
		fmt.Fprintln(os.Stderr, "Socket", socketPath, "does not exist")
		os.Exit(1)
	}

	// Get original socket permissions and ownership
	info, err := os.Stat(socketPath)
	if err != nil {
		fmt.Fprintln(os.Stderr, "Failed to stat socket:", err)
		os.Exit(1)
	}
	mode := info.Mode()

	// Get owner/group (Unix specific)
	stat := info.Sys().(*syscall.Stat_t)
	uid := stat.Uid
	gid := stat.Gid

	// Rename the original socket
	if err := os.Rename(socketPath, origSocketPath); err != nil {
		fmt.Fprintln(os.Stderr, "Failed to rename socket:", err)
		os.Exit(1)
	}

	// Cleanup function to restore the original socket
	cleanup := func() {
		fmt.Fprintln(os.Stderr, "Shutting down, restoring original socket...")
		if _, err := os.Stat(socketPath); err == nil {
			os.Remove(socketPath)
		}
		if _, err := os.Stat(origSocketPath); err == nil {
			if err := os.Rename(origSocketPath, socketPath); err != nil {
				fmt.Fprintln(os.Stderr, "Warning: Failed to restore original socket:", err)
			} else {
				fmt.Fprintln(os.Stderr, "Original socket restored to", socketPath)
			}
		}
	}

	// Create new listener on the original path
	listener, err := net.Listen("unix", socketPath)
	if err != nil {
		// Try to restore the original socket on failure
		cleanup()
		fmt.Fprintln(os.Stderr, "Failed to create listener:", err)
		os.Exit(1)
	}

	// Set permissions and ownership on the new socket
	if err := os.Chmod(socketPath, mode); err != nil {
		fmt.Fprintln(os.Stderr, "Warning: Failed to set socket permissions:", err)
	}
	if err := os.Chown(socketPath, int(uid), int(gid)); err != nil {
		fmt.Fprintln(os.Stderr, "Warning: Failed to set socket ownership:", err)
	}

	fmt.Fprintf(os.Stderr, "Listening on %s, forwarding to %s\n", socketPath, origSocketPath)
	fmt.Fprintln(os.Stderr, "Press Ctrl+C to stop and restore the original socket")

	// Flag to track if we're shutting down
	shutdownFlag := false

	// Handle signals in a goroutine
	go func() {
		<-sigChan
		shutdownFlag = true

		fmt.Fprintln(os.Stderr, "\nReceived interrupt signal")
		listener.Close()
	}()

	for connID := 0; ; connID++ {
		client, err := listener.Accept()
		if err != nil {

			if shutdownFlag {
				break
			}
			fmt.Fprintln(os.Stderr, "Accept error:", err)
			continue
		}

		go handleConnection(client, origSocketPath, connID)
	}

	cleanup()
}

func handleConnection(client net.Conn, origSocketPath string, connID int) {
	defer client.Close()

	fmt.Fprintf(os.Stderr, "[Connection #%d] New connection\n", connID)

	// Connect to the original socket
	server, err := net.Dial("unix", origSocketPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "[Connection #%d] Failed to connect to original socket: %v\n", connID, err)
		return
	}
	defer server.Close()

	var wg sync.WaitGroup
	wg.Add(2)

	// Client -> Server
	go func() {
		defer wg.Done()
		copyAndLog(server, client, connID, "CLIENT->SERVER")
	}()

	// Server -> Client
	go func() {
		defer wg.Done()
		copyAndLog(client, server, connID, "SERVER->CLIENT")
	}()

	wg.Wait()
	fmt.Fprintf(os.Stderr, "[Connection #%d] Connection closed\n", connID)
}

func copyAndLog(dst io.Writer, src io.Reader, connID int, direction string) {
	tee := io.TeeReader(src, os.Stdout)

	_, err := io.Copy(dst, tee)

	if err != nil && err != io.EOF {
		fmt.Fprintf(os.Stderr, "[Connection #%d] %s error: %v\n", connID, direction, err)
	}
}