└── listenbuddy.go


/listenbuddy.go:
--------------------------------------------------------------------------------
  1 | package main
  2 | 
  3 | import (
  4 | 	"flag"
  5 | 	"fmt"
  6 | 	"io"
  7 | 	"log"
  8 | 	"net"
  9 | 	"os"
 10 | 	"os/signal"
 11 | 	"strings"
 12 | 	"sync"
 13 | 	"syscall"
 14 | )
 15 | 
 16 | var listen = flag.String("listen", "", "port (and optionally address) to listen on")
 17 | var speak = flag.String("speak", "", "address and port to connect to")
 18 | 
 19 | var connections = make(map[*net.TCPConn]struct{}, 100)
 20 | var cMu sync.Mutex
 21 | 
 22 | func main() {
 23 | 	flag.Parse()
 24 | 	if *listen == "" || *speak == "" {
 25 | 		fmt.Println(`
 26 | Forwards all connections from a given port to a different address and port.
 27 | Example:
 28 | 	listenbuddy -listen :8000 -speak localhost:80
 29 | `)
 30 | 		flag.Usage()
 31 | 		return
 32 | 	}
 33 | 	log.SetPrefix("listenbuddy ")
 34 | 
 35 | 	speakAddr, err := net.ResolveTCPAddr("tcp", *speak)
 36 | 	if err != nil {
 37 | 		log.Println(err)
 38 | 		return
 39 | 	}
 40 | 
 41 | 	listenAddr, err := net.ResolveTCPAddr("tcp", *listen)
 42 | 	if err != nil {
 43 | 		log.Println(err)
 44 | 		return
 45 | 	}
 46 | 	ln, err := net.ListenTCP("tcp", listenAddr)
 47 | 	if err != nil {
 48 | 		log.Println(err)
 49 | 		return
 50 | 	}
 51 | 	go handleSignals()
 52 | 	for {
 53 | 		conn, err := ln.AcceptTCP()
 54 | 		if err != nil {
 55 | 			log.Println("accept", err)
 56 | 			return
 57 | 		}
 58 | 		go handleConn(speakAddr, conn)
 59 | 	}
 60 | }
 61 | 
 62 | func handleSignals() {
 63 | 	ch := make(chan os.Signal, 1)
 64 | 	signal.Notify(ch, syscall.SIGUSR1)
 65 | 	for _ = range ch {
 66 | 		log.Println("closing")
 67 | 		closeAllConnections()
 68 | 	}
 69 | }
 70 | 
 71 | func closeAllConnections() {
 72 | 	cMu.Lock()
 73 | 	defer cMu.Unlock()
 74 | 	for c, _ := range connections {
 75 | 		c.CloseWrite()
 76 | 	}
 77 | }
 78 | 
 79 | func addConnection(c *net.TCPConn) {
 80 | 	cMu.Lock()
 81 | 	connections[c] = struct{}{}
 82 | 	cMu.Unlock()
 83 | }
 84 | 
 85 | func removeConnection(c *net.TCPConn) {
 86 | 	cMu.Lock()
 87 | 	delete(connections, c)
 88 | 	cMu.Unlock()
 89 | }
 90 | 
 91 | func copyConn(dst, src *net.TCPConn) {
 92 | 	addConnection(src)
 93 | 	_, err := io.Copy(dst, src)
 94 | 	if err != nil {
 95 | 		// We commonly get use of closed network connection when a server shuts down
 96 | 		// active
 97 | 		if !strings.Contains(err.Error(), "use of closed network connection") {
 98 | 			log.Println(err)
 99 | 		}
100 | 	}
101 | 	src.Close()
102 | 	dst.CloseWrite()
103 | 	removeConnection(src)
104 | }
105 | 
106 | // Any time we get an inbound connection, connect to the "speak" host and port,
107 | // and spawn two goroutines: one to copy data in each direction.
108 | // When either connection generates an error (terminating the Copy call), close
109 | // both connections.
110 | func handleConn(speakAddr *net.TCPAddr, hearing *net.TCPConn) {
111 | 	speaking, err := net.DialTCP("tcp", nil, speakAddr)
112 | 	if err != nil {
113 | 		log.Println(err)
114 | 		hearing.Close()
115 | 		return
116 | 	}
117 | 	go copyConn(speaking, hearing)
118 | 	copyConn(hearing, speaking)
119 | }
120 | 


--------------------------------------------------------------------------------