Browse Source

Make the address argument work, support AF_UNIX.

Yawning Angel 1 year ago
parent
commit
12bcc7adb0
1 changed files with 46 additions and 18 deletions
  1. 46 18
      main.go

+ 46 - 18
main.go

@@ -32,37 +32,65 @@ import (
 	"io"
 	"log"
 	"net"
+	"os"
+	"os/signal"
+	"strings"
 	"sync"
 
 	"git.schwanenlied.me/yawning/shittysocks.git/socks5"
 )
 
 func main() {
-	addr := flag.String("address", "127.0.0.1:1080", "Address to listen on.")
+	const (
+		netUnix    = "unix"
+		prefixUnix = netUnix + ":"
+	)
+
+	pAddr := flag.String("address", "127.0.0.1:1080", "Address to listen on.")
+	flag.Parse()
+
+	addr, netwerk := *pAddr, "tcp"
+	if strings.HasPrefix(addr, prefixUnix) {
+		netwerk = netUnix
+		addr = strings.TrimPrefix(addr, prefixUnix)
+	}
 
-	ln, err := net.Listen("tcp", *addr)
+	ln, err := net.Listen(netwerk, addr)
 	if err != nil {
 		log.Fatalf("Failed to listen: %v", err)
 	}
 	log.Printf("Listening on: %v", ln.Addr())
-	defer ln.Close()
-
-	id := 0
-	for {
-		conn, err := ln.Accept()
-		if err != nil {
-			if e, ok := err.(net.Error); ok && !e.Temporary() {
-				log.Printf("Critical accept failure: %v", err)
-				return
+	defer func() {
+		if netwerk == netUnix {
+			os.Remove(addr)
+		}
+		ln.Close()
+	}()
+
+	ch := make(chan os.Signal, 1)
+	signal.Notify(ch, os.Interrupt)
+
+	go func() {
+		id := 0
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				if e, ok := err.(net.Error); ok && !e.Temporary() {
+					log.Printf("Critical accept failure: %v", err)
+					return
+				}
+				continue
 			}
-			continue
+
+			rAddr := conn.RemoteAddr()
+			log.Printf("Accepted new connection: %v (%d)", rAddr, id)
+			go connWorker(conn, id)
+			id++
 		}
+	}()
 
-		rAddr := conn.RemoteAddr()
-		log.Printf("Accepted new connection: %v (%d)", rAddr, id)
-		go connWorker(conn, id)
-		id++
-	}
+	<-ch
+	log.Printf("Terminating.")
 }
 
 func connWorker(conn net.Conn, id int) {
@@ -74,7 +102,7 @@ func connWorker(conn net.Conn, id int) {
 		return
 	}
 
-	log.Printf("%d: Target: %v (Auth: %v:%v)", id, req.Target, req.Uname, req.Passwd)
+	log.Printf("%d: Target: %v (Auth: `%v`/`%v`)", id, req.Target, string(req.Uname), string(req.Passwd))
 	upConn, err := net.Dial("tcp", req.Target)
 	if err != nil {
 		req.Reply(socks5.ErrorToReplyCode(err))