creator-hub/cmd/control-plane/network.go

package main

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"git.ipao.vip/rogee/creator-hub/internal/hub"
)

const networkExitObservationURL = "https://ipinfo.io/json"

type networkExitProbe interface {
	Check(context.Context, hub.NetworkExitAccess) (hub.ExitObservation, string)
}

type httpNetworkExitProbe struct {
	endpoint string
	client   *http.Client
}

func defaultNetworkExitProbe() networkExitProbe {
	return httpNetworkExitProbe{endpoint: networkExitObservationURL, client: &http.Client{Timeout: 20 * time.Second}}
}

func (probe httpNetworkExitProbe) Check(ctx context.Context, exit hub.NetworkExitAccess) (hub.ExitObservation, string) {
	proxyURL := &url.URL{Scheme: exit.Protocol, Host: net.JoinHostPort(exit.Host, fmt.Sprint(exit.Port))}
	proxyUsername := exit.Username
	if exit.Username != "" {
		proxyURL.User = url.UserPassword(exit.Username, exit.Password)
	}
	transport := http.DefaultTransport.(*http.Transport).Clone()
	if configured, ok := probe.client.Transport.(*http.Transport); ok {
		transport = configured.Clone()
	}
	transport.Proxy = http.ProxyURL(proxyURL)
	if exit.Protocol == "socks4" {
		transport.Proxy = nil
		transport.DialContext = socks4DialContext(proxyURL.Host, proxyUsername)
	}
	defer transport.CloseIdleConnections()
	client := *probe.client
	client.Transport = transport
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, probe.endpoint, nil)
	if err != nil {
		return hub.ExitObservation{}, "proxy_check_failed"
	}
	response, err := client.Do(request)
	if err != nil {
		if strings.Contains(strings.ToLower(err.Error()), "auth") {
			return hub.ExitObservation{}, "proxy_auth_failed"
		}
		return hub.ExitObservation{}, "proxy_check_failed"
	}
	defer response.Body.Close()
	if response.StatusCode == http.StatusProxyAuthRequired {
		return hub.ExitObservation{}, "proxy_auth_failed"
	}
	if response.StatusCode != http.StatusOK {
		return hub.ExitObservation{}, "proxy_check_failed"
	}
	var observed struct {
		IP     string `json:"ip"`
		Region string `json:"region"`
	}
	decoder := json.NewDecoder(io.LimitReader(response.Body, 64<<10))
	if err := decoder.Decode(&observed); err != nil || net.ParseIP(observed.IP) == nil || len(observed.Region) > 64 {
		return hub.ExitObservation{}, "exit_observation_invalid"
	}
	return hub.ExitObservation{PublicIP: observed.IP, Region: observed.Region}, ""
}

func socks4DialContext(proxyAddress, userID string) func(context.Context, string, string) (net.Conn, error) {
	return func(ctx context.Context, network, address string) (net.Conn, error) {
		connection, err := (&net.Dialer{}).DialContext(ctx, network, proxyAddress)
		if err != nil {
			return nil, err
		}
		failed := true
		defer func() {
			if failed {
				_ = connection.Close()
			}
		}()
		host, portText, err := net.SplitHostPort(address)
		if err != nil {
			return nil, errors.New("invalid SOCKS4 destination")
		}
		port, err := net.LookupPort("tcp", portText)
		if err != nil {
			return nil, errors.New("invalid SOCKS4 destination port")
		}
		request := []byte{4, 1, 0, 0, 0, 0, 0, 1}
		binary.BigEndian.PutUint16(request[2:4], uint16(port))
		if ip := net.ParseIP(host).To4(); ip != nil {
			copy(request[4:8], ip)
		}
		request = append(request, userID...)
		request = append(request, 0)
		if net.ParseIP(host).To4() == nil {
			request = append(request, host...)
			request = append(request, 0)
		}
		if deadline, ok := ctx.Deadline(); ok {
			_ = connection.SetDeadline(deadline)
		}
		if _, err := connection.Write(request); err != nil {
			return nil, err
		}
		response := make([]byte, 8)
		if _, err := io.ReadFull(connection, response); err != nil || response[1] != 90 {
			return nil, errors.New("SOCKS4 proxy rejected connection")
		}
		_ = connection.SetDeadline(time.Time{})
		failed = false
		return connection, nil
	}
}

type gatewayNetworkExit struct {
	Protocol string `json:"protocol"`
	Host     string `json:"host"`
	Port     int    `json:"port"`
	Username string `json:"username,omitempty"`
	Password string `json:"password,omitempty"`
}

func gatewayNetworkExitFor(exit hub.NetworkExitAccess) gatewayNetworkExit {
	return gatewayNetworkExit{
		Protocol: exit.Protocol,
		Host:     exit.Host,
		Port:     exit.Port,
		Username: exit.Username,
		Password: exit.Password,
	}
}