UI improvements and Windows color support

Author: rtpt-erikgeiser

README.md

@@ -44,7 +44,13 @@ $ curl --proxy 'socks5://127.0.0.1:1080' 'http://10.0.0.2'
 ```
 
 You can also generate a connection key with `resocks generate` and pass it to
-the `listen` command to avoid generating a new connection key every time.
+the `listen` command to avoid generating a new connection key every time. It
+can also be specified via an environment variable:
+
+```bash
+$ export RESOCKS_KEY="$(resocks generate)"
+$ resocks listen
+```
 
 ## Security
 
@@ -117,3 +123,10 @@ directly into the binary, use the following command:
 go run . generate  # generate a connection key
 go build -ldflags="-X main.defaultConnectionKey=YOUR_CONNECTION_KEY"
 ```
+
+Similarly, the default connect back address can also be statically compiled into
+the binary:
+
+```bash
+go build -ldflags="-X main.defaultConnectBackAddress=192.0.2.1"
+```

main.go

@@ -26,7 +26,10 @@ const (
 	ConnectionKeyEnvVariable = "RESOCKS_KEY"
 )
 
-var defaultConnectionKey = ""
+var (
+	defaultConnectionKey      = ""
+	defaultConnectBackAddress = ""
+)
 
 func main() {
 	err := run()

relay.go

@@ -19,15 +19,25 @@ func relayCommand() *cobra.Command {
 	connectionKey := fromEnvWithFallback(ConnectionKeyEnvVariable, defaultConnectionKey)
 	insecure := false
 
+	cobraArgs := cobra.ExactArgs(1)
+	if defaultConnectBackAddress != "" {
+		cobraArgs = cobra.MaximumNArgs(1)
+	}
+
 	relayCmd := &cobra.Command{
 		Use:           fmt.Sprintf("%s <connect back address> --key <connection key>", binaryName()),
 		Short:         fmt.Sprintf("Connect back to an %s listener and relay the SOCKS5 traffic", binaryName()),
 		SilenceErrors: true,
 		SilenceUsage:  true,
-		Args:          cobra.ExactArgs(1),
+		Args:          cobraArgs,
 		RunE: func(cmd *cobra.Command, args []string) error {
+			connectBackAddress := defaultConnectBackAddress
+			if len(args) > 0 {
+				connectBackAddress = args[0]
+			}
+
 			return runRemoteProxyRelay(
-				withDefaultPort(args[0], DefaultListenPort), // connect back address
+				withDefaultPort(connectBackAddress, DefaultListenPort),
 				connectionKey,
 				timeout,
 				reconnectAfter,

ui.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"strings"
+	"sync"
 	"time"
 
 	"resocks/kbtls"
@@ -40,6 +41,8 @@ type model struct {
 	socksAddress  string
 
 	noColor bool
+
+	lock sync.Mutex
 }
 
 var _ tea.Model = &model{}
@@ -57,7 +60,9 @@ func startUI(
 	})
 
 	if !noColor {
-		err := enableVirtualTerminalProcessing()
+		resetTerminal, err := prepareTerminal()
+		defer resetTerminal()
+
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "Warning: Could not enable virtual terminal processing: %v, "+
 				"disabling colored output\n", err)
@@ -96,6 +101,9 @@ func (m *model) Init() tea.Cmd {
 }
 
 func (m *model) Update(message tea.Msg) (tea.Model, tea.Cmd) {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
 	switch msg := message.(type) {
 	case proxyrelay.Event:
 		switch msg.Type {
@@ -128,6 +136,11 @@ func (m *model) Update(message tea.Msg) (tea.Model, tea.Cmd) {
 	case tea.KeyMsg:
 		if msg.String() == "ctrl+c" {
 			m.quitting = true
+			if m.connection != nil {
+				m.connection.End = time.Now()
+				m.connection.Error = "shutdown"
+				m.connection = nil
+			}
 
 			return m, tea.Quit
 		}
@@ -181,13 +194,20 @@ func (m *model) errorsView() string {
 }
 
 func (m *model) listenerConfigView() string {
-	listeningLineStyle := bold
+	address := m.listenAddress
+	if strings.HasPrefix(address, ":") {
+		address = "*" + address
+	}
+
+	addressStyle := m.style(brightMagenta, bold)
+
 	if m.socksActive {
-		listeningLineStyle = dim
+		addressStyle = m.style(dim)
+		address += " (inactive)"
 	}
 
-	view := m.style(listeningLineStyle) + "Listening On  : " + m.style() +
-		m.style(brightMagenta, listeningLineStyle) + m.listenAddress + m.style()
+	view := m.style(bold) + "Listening On  : " + m.style() +
+		addressStyle + address + m.style()
 	view += m.style(bold) + "\nConnection Key: " + m.style() + m.style(brightMagenta, bold) + m.connectionKey + m.style()
 
 	if m.insecure {

ui_default.go

@@ -2,6 +2,6 @@
 
 package main
 
-func enableVirtualTerminalProcessing() error {
-	return nil // we only have to do this on Windows
+func prepareTerminal() (reset func(), err error) {
+	return func() {}, nil // we only have to do this on Windows
 }

ui_windows.go

@@ -9,32 +9,39 @@ import (
 	"golang.org/x/sys/windows"
 )
 
-func enableVirtualTerminalProcessing() error {
+func prepareTerminal() (reset func(), err error) {
 	var mode uint32
 
 	stdoutHandle := windows.Handle(os.Stdout.Fd())
 
-	err := windows.GetConsoleMode(stdoutHandle, &mode)
+	err = windows.GetConsoleMode(stdoutHandle, &mode)
 	if err != nil {
-		return fmt.Errorf("stdout: get console mode: %w", err)
+		return nil, fmt.Errorf("stdout: get console mode: %w", err)
 	}
 
 	err = windows.SetConsoleMode(stdoutHandle, mode|windows.ENABLE_VIRTUAL_TERMINAL_PROCESSING)
 	if err != nil {
-		return fmt.Errorf("stdout: set console mode: %w", err)
+		return nil, fmt.Errorf("stdout: set console mode: %w", err)
 	}
 
 	stderrHandle := windows.Handle(os.Stderr.Fd())
 
 	err = windows.GetConsoleMode(stderrHandle, &mode)
 	if err != nil {
-		return fmt.Errorf("stderr: get console mode: %w", err)
+		return func() {
+			_ = windows.SetConsoleMode(stdoutHandle, mode)
+		}, fmt.Errorf("stderr: get console mode: %w", err)
 	}
 
 	err = windows.SetConsoleMode(stderrHandle, mode|windows.ENABLE_VIRTUAL_TERMINAL_PROCESSING)
 	if err != nil {
-		return fmt.Errorf("stderr: set console mode: %w", err)
+		return func() {
+			_ = windows.SetConsoleMode(stdoutHandle, mode)
+		}, fmt.Errorf("stderr: set console mode: %w", err)
 	}
 
-	return nil
+	return func() {
+		_ = windows.SetConsoleMode(stdoutHandle, mode)
+		_ = windows.SetConsoleMode(stderrHandle, mode)
+	}, nil
 }