Merge pull request #21 from projectdiscovery/maintenance-freeport-bind

bind logic refactor

Author: Mzack9999

go.mod

@@ -3,6 +3,7 @@ module github.com/projectdiscovery/simplehttpserver
 go 1.15
 
 require (
+	github.com/phayes/freeport v0.0.0-20180830031419-95f893ade6f2
 	github.com/projectdiscovery/gologger v1.1.4
 	gopkg.in/yaml.v2 v2.4.0
 )

go.sum

@@ -19,6 +19,8 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJ
 github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/modern-go/reflect2 v1.0.1 h1:9f412s+6RmYXLWZSEzVVgPGK7C2PphHj5RJrvfx9AWI=
 github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
+github.com/phayes/freeport v0.0.0-20180830031419-95f893ade6f2 h1:JhzVVoYvbOACxoUmOs6V/G4D5nPVUW73rKvXxP4XUJc=
+github.com/phayes/freeport v0.0.0-20180830031419-95f893ade6f2/go.mod h1:iIss55rKnNBTvrwdmkUpLnDpZoAHvWaiq5+iMmen4AE=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/projectdiscovery/gologger v1.1.4 h1:qWxGUq7ukHWT849uGPkagPKF3yBPYAsTtMKunQ8O2VI=

internal/runner/runner.go

@@ -1,6 +1,8 @@
 package runner
 
 import (
+	"github.com/projectdiscovery/gologger"
+	"github.com/projectdiscovery/simplehttpserver/pkg/binder"
 	"github.com/projectdiscovery/simplehttpserver/pkg/httpserver"
 	"github.com/projectdiscovery/simplehttpserver/pkg/tcpserver"
 )
@@ -14,6 +16,16 @@ type Runner struct {
 
 func New(options *Options) (*Runner, error) {
 	r := Runner{options: options}
+	// Check if the process can listen on the specified ip:port
+	if !binder.CanListenOn(r.options.ListenAddress) {
+		newListenAddress, err := binder.GetRandomListenAddress(r.options.ListenAddress)
+		if err != nil {
+			return nil, err
+		}
+		gologger.Print().Msgf("Can't listen on %s: %s - Using %s\n", r.options.ListenAddress, err, newListenAddress)
+		r.options.ListenAddress = newListenAddress
+	}
+
 	if r.options.EnableTCP {
 		serverTCP, err := tcpserver.New(tcpserver.Options{
 			Listen:  r.options.ListenAddress,
@@ -55,16 +67,25 @@ func New(options *Options) (*Runner, error) {
 
 func (r *Runner) Run() error {
 	if r.options.EnableTCP {
+		gologger.Print().Msgf("Serving TCP rule based server on tcp://%s", r.options.ListenAddress)
 		return r.serverTCP.ListenAndServe()
 	}
 
 	if r.options.HTTPS {
+		gologger.Print().Msgf("Serving %s on https://%s/...", r.options.Folder, r.options.ListenAddress)
 		return r.httpServer.ListenAndServeTLS()
 	}
 
+	gologger.Print().Msgf("Serving %s on http://%s/...", r.options.Folder, r.options.ListenAddress)
 	return r.httpServer.ListenAndServe()
 }
 
 func (r *Runner) Close() error {
+	if r.serverTCP != nil {
+		r.serverTCP.Close()
+	}
+	if r.httpServer != nil {
+		r.httpServer.Close()
+	}
 	return nil
 }

pkg/binder/binder.go

@@ -0,0 +1,32 @@
+package binder
+
+import (
+	"fmt"
+	"net"
+
+	"github.com/phayes/freeport"
+)
+
+func CanListenOn(address string) bool {
+	listener, err := net.Listen("tcp4", address)
+	if err != nil {
+		return false
+	}
+	defer listener.Close()
+	return true
+}
+
+func GetRandomListenAddress(currentAddress string) (string, error) {
+	addrOrig, _, err := net.SplitHostPort(currentAddress)
+	if err != nil {
+		return "", err
+	}
+
+	newPort, err := freeport.GetFreePort()
+	if err != nil {
+		return "", err
+	}
+
+	return net.JoinHostPort(addrOrig, fmt.Sprintf("%d", newPort)), nil
+
+}

pkg/httpserver/httpserver.go

@@ -1,15 +1,9 @@
 package httpserver
 
 import (
-	"errors"
 	"net"
 	"net/http"
-	"os"
-	"runtime"
-	"strconv"
-	"syscall"
 
-	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/simplehttpserver/pkg/sslcert"
 )
 
@@ -46,26 +40,10 @@ func New(options *Options) (*HTTPServer, error) {
 }
 
 func (t *HTTPServer) ListenAndServe() error {
-	var err error
-retry_listen:
-	gologger.Print().Msgf("Serving %s on http://%s/...", t.options.Folder, t.options.ListenAddress)
-	err = http.ListenAndServe(t.options.ListenAddress, t.layers)
-	if err != nil {
-		if isErrorAddressAlreadyInUse(err) {
-			gologger.Print().Msgf("Can't listen on %s: %s - retrying with another port\n", t.options.ListenAddress, err)
-			newListenAddress, err := incPort(t.options.ListenAddress)
-			if err != nil {
-				return err
-			}
-			t.options.ListenAddress = newListenAddress
-			goto retry_listen
-		}
-	}
-	return nil
+	return http.ListenAndServe(t.options.ListenAddress, t.layers)
 }
 
 func (t *HTTPServer) ListenAndServeTLS() error {
-	gologger.Print().Msgf("Serving %s on https://%s/...", t.options.Folder, t.options.ListenAddress)
 	if t.options.Certificate == "" || t.options.CertificateKey == "" {
 		tlsOptions := sslcert.DefaultOptions
 		tlsOptions.Host = t.options.CertificateDomain
@@ -83,38 +61,6 @@ func (t *HTTPServer) ListenAndServeTLS() error {
 	return http.ListenAndServeTLS(t.options.ListenAddress, t.options.Certificate, t.options.CertificateKey, t.layers)
 }
 
-func isErrorAddressAlreadyInUse(err error) bool {
-	var eOsSyscall *os.SyscallError
-	if !errors.As(err, &eOsSyscall) {
-		return false
-	}
-	var errErrno syscall.Errno // doesn't need a "*" (ptr) because it's already a ptr (uintptr)
-	if !errors.As(eOsSyscall, &errErrno) {
-		return false
-	}
-	if errErrno == syscall.EADDRINUSE {
-		return true
-	}
-	const WSAEADDRINUSE = 10048
-	if runtime.GOOS == "windows" && errErrno == WSAEADDRINUSE {
-		return true
-	}
-	return false
-}
-
-func incPort(address string) (string, error) {
-	addrOrig, portOrig, err := net.SplitHostPort(address)
-	if err != nil {
-		return address, err
-	}
-
-	// increment port
-	portNumber, err := strconv.Atoi(portOrig)
-	if err != nil {
-		return address, err
-	}
-	portNumber++
-	newPort := strconv.FormatInt(int64(portNumber), 10)
-
-	return net.JoinHostPort(addrOrig, newPort), nil
+func (t *HTTPServer) Close() error {
+	return nil
 }

pkg/tcpserver/tcpserver.go

@@ -36,7 +36,6 @@ func (t *TCPServer) AddRule(rule Rule) error {
 }
 
 func (t *TCPServer) ListenAndServe() error {
-	gologger.Print().Msgf("Serving %s on tcp://%s", t.options.Listen)
 	listener, err := net.Listen("tcp4", t.options.Listen)
 	if err != nil {
 		return err