From e46ad13c67b6b86f97a1234e5aa14b20e3e23992 Mon Sep 17 00:00:00 2001 From: onionj98 Date: Sat, 9 Mar 2024 15:44:34 +0330 Subject: [PATCH] feat: :zap: add workers --- .gitignore | 3 +- Makefile | 2 +- go.mod | 2 +- main.go | 155 ++++++++++++++++++++++++++--------------------------- 4 files changed, 79 insertions(+), 83 deletions(-) diff --git a/.gitignore b/.gitignore index a2f1740..c7a9176 100644 --- a/.gitignore +++ b/.gitignore @@ -21,4 +21,5 @@ go.work yourip* build -ipGeolocationcidr.json \ No newline at end of file +ipGeolocationcidr.json +.DS_Store \ No newline at end of file diff --git a/Makefile b/Makefile index ac79fe8..9f6845b 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ # Define the output directory for the builds OUTPUT_DIR := ./build -version := 1.7.0 +version := 1.8.0 # Define the name of your binary BINARY_NAME := yourip diff --git a/go.mod b/go.mod index ceb4d87..df916e1 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,3 @@ module yourip -go 1.19 +go 1.21 diff --git a/main.go b/main.go index 4d144f4..52202ce 100644 --- a/main.go +++ b/main.go @@ -9,6 +9,7 @@ import ( "encoding/json" "flag" "fmt" + "io" "math/rand" "net" "os" @@ -47,6 +48,7 @@ type JsonResponse struct { var ( bindAddress string + workers int useIPGeolocation bool ipGeolocation *geolocation.IPGeolocation ) @@ -55,128 +57,121 @@ func main() { flag.StringVar(&bindAddress, "bind", ":8080", "The address to bind the TCP server to. (shorthand -b)") flag.StringVar(&bindAddress, "b", ":8080", "(shorthand for --bind)") + flag.IntVar(&workers, "worker", 3, "number of workers") + flag.IntVar(&workers, "w", 3, "number of workers") + flag.BoolVar(&useIPGeolocation, "geolocation", false, "ip geolocation service (shorthand -g)") flag.BoolVar(&useIPGeolocation, "g", false, "(shorthand for --geolocation)") flag.Parse() if useIPGeolocation { - ipGeolocation = geolocation.New(time.Duration(48)) + ipGeolocation = geolocation.New(time.Duration(100)) } tcpAddr, err := net.ResolveTCPAddr("tcp4", bindAddress) - checkError(err) + if err != nil { + fmt.Fprintf(os.Stderr, "Fatal error: %s", err.Error()) + os.Exit(1) + } listener, err := net.ListenTCP("tcp", tcpAddr) - checkError(err) + if err != nil { + fmt.Fprintf(os.Stderr, "Fatal error: %s", err.Error()) + os.Exit(1) + } + + defer listener.Close() fmt.Printf("Start a new listener on %s\nversion:%s\n", tcpAddr.String(), version) - for { - if conn, err := listener.Accept(); err == nil { - go Handler(conn) - } + for ; workers != 0; workers-- { + fmt.Println("Runing Worker", workers) + go Worker(listener) } + + ch := make(chan struct{}) + <-ch } -// Accept the connection, Prepare Response and send it to client -func Handler(conn net.Conn) { - defer func() { - conn.Close() - fmt.Printf("(%s) %s Closed\n", time.Now().String()[:23], conn.RemoteAddr().String()) - }() +func Worker(listener *net.TCPListener) { + for { + conn, err := listener.Accept() + if err != nil { + fmt.Println("Accept err:", err) + continue + } - fmt.Printf("(%s) %s Accepted\n", time.Now().String()[:23], conn.RemoteAddr().String()) + fmt.Printf("(%s) %s Accepted\n", time.Now().String()[:23], conn.RemoteAddr().String()) - conn.SetDeadline(time.Now().Add(time.Minute)) + err = conn.SetDeadline(time.Now().Add(time.Minute)) + if err != nil { + fmt.Println("SetDeadline err:", err) + continue + } - // Read request from tcp connection - RequestBuffer := make([]byte, 1024) - numberOfBytes, err := conn.Read(RequestBuffer) - if err == nil { - fmt.Printf("(%s) Read %d byte from %s\n", time.Now().String()[:23], numberOfBytes, conn.RemoteAddr().String()) + go handleConnection(conn) } +} - response := "" - animationMode := "" - - // Chose response type - if numberOfBytes <= 1 { - response, _ = CreateResponse(conn, ResponseModeText) - - } else { - request := string(RequestBuffer[:numberOfBytes]) - - if strings.Contains(request, "GET /json") { - response, err = CreateResponse(conn, ResponseModeHTTPJson) - if err != nil { - fmt.Println(err.Error()) - return - } - } else if request == AnimationModeBanner || request == AnimationModeFlight { - response, _ = CreateResponse(conn, ResponseModeText) - animationMode = request +func handleConnection(conn net.Conn) { + defer func() { + err := conn.Close() + if err != nil { + fmt.Printf("(%s) %s Close err: %s\n", time.Now().String()[:23], conn.RemoteAddr().String(), err.Error()) } else { - response, _ = CreateResponse(conn, ResponseModeHTTPText) - } - } - - // Write response to TCP connection - if animationMode != "" { - StreamAnimation(conn, response, animationMode) - } else { - numberOfBytes, err = conn.Write([]byte(response)) - if err == nil { - fmt.Printf("(%s) Write %d byte to %s\n", time.Now().String()[:23], numberOfBytes, conn.RemoteAddr().String()) + fmt.Printf("(%s) %s Closed\n", time.Now().String()[:23], conn.RemoteAddr().String()) } - } - -} + }() -func checkError(err error) { + requestBuffer := make([]byte, 1024) + bytesRead, err := conn.Read(requestBuffer) if err != nil { - fmt.Fprintf(os.Stderr, "Fatal error: %s", err.Error()) - os.Exit(1) + if err == io.EOF { + fmt.Println("Connection closed by client") + } else { + fmt.Println("Error reading request:", err) + return + } } -} -// Return Client IP in these formats: -// -// simple text: for pure TCP like NetCat -// http text: for browsers and curl -// http JSON: for API Call -func CreateResponse(conn net.Conn, mode ResponseMode) (string, error) { remoteAddr := conn.RemoteAddr().(*net.TCPAddr) + remoteAddrStr := remoteAddr.IP.String() country := "" - separator := "" - if useIPGeolocation { country, _ = ipGeolocation.Query(remoteAddr.IP) - separator = " " } - switch mode { + // for netcat client: + if bytesRead <= 1 { + conn.Write([]byte(fmt.Sprintf("%s %s", remoteAddrStr, country))) + return + } - case ResponseModeHTTPText: - response := remoteAddr.IP.String() + separator + country - return fmt.Sprintf(RESPONSE, len(response), "text/plain", response), nil + request := string(requestBuffer[:bytesRead]) - case ResponseModeHTTPJson: - jsonResponse := JsonResponse{IP: remoteAddr.IP.String(), Country: country} - jsonResponseByte, err := json.Marshal(jsonResponse) + if strings.Contains(request, " /json") { + jsonResponseByte, err := json.Marshal(JsonResponse{IP: remoteAddrStr, Country: country}) if err != nil { - return "", fmt.Errorf("failed to marshal message to JSON: %v", err) + fmt.Println("create json response err:", err) + return + } + conn.Write([]byte(fmt.Sprintf(RESPONSE, len(jsonResponseByte), "application/json", jsonResponseByte))) + return + + } else { + ipWithCountry := fmt.Sprintf("%s %s", remoteAddrStr, country) + + if request == AnimationModeBanner || request == AnimationModeFlight { + StreamAnimation(conn, ipWithCountry, request) + return } - return fmt.Sprintf(RESPONSE, len(jsonResponseByte), "application/json", jsonResponseByte), nil - default: - return remoteAddr.IP.String() + separator + country, nil + conn.Write([]byte(fmt.Sprintf(RESPONSE, len(ipWithCountry), "text/plain", ipWithCountry))) } } -// Stream Client IP Like an Animation! -// -// > for example client open a tcp connection with netcat and send `1` +// Stream Client IP Like an Animation (for example client open a tcp connection with netcat and send `1`) func StreamAnimation(conn net.Conn, response string, animationMode string) { switch animationMode { case AnimationModeBanner: