chis_osi/osi/transport.go

package osi

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"golang.org/x/net/proxy"
)

type TransportConfig struct {
	Timeout     time.Duration
	Socks5Proxy string
}

type Transport struct {
	timeout time.Duration
	dialer  proxy.Dialer
}

func NewTransport(config TransportConfig) (*Transport, error) {
	timeout := config.Timeout
	if timeout <= 0 {
		timeout = 20 * time.Second
	}

	dialer := proxy.Dialer(proxy.Direct)
	if strings.TrimSpace(config.Socks5Proxy) != "" {
		proxyDialer, err := socks5Dialer(config.Socks5Proxy)
		if err != nil {
			return nil, err
		}
		dialer = proxyDialer
	}

	return &Transport{timeout: timeout, dialer: dialer}, nil
}

func (t *Transport) PostJSON(ctx context.Context, targetURL string, payload any, headers map[string]string) (int, []byte, error) {
	body, err := json.Marshal(payload)
	if err != nil {
		return 0, nil, fmt.Errorf("marshal json payload: %w", err)
	}

	target, err := url.Parse(targetURL)
	if err != nil {
		return 0, nil, fmt.Errorf("parse target url: %w", err)
	}
	if target.Scheme != "http" {
		return 0, nil, fmt.Errorf("unsupported target scheme %q", target.Scheme)
	}

	reqCtx := ctx
	cancel := func() {}
	if _, ok := ctx.Deadline(); !ok && t.timeout > 0 {
		reqCtx, cancel = context.WithTimeout(ctx, t.timeout)
	}
	defer cancel()

	conn, err := dialWithContext(reqCtx, t.dialer, "tcp", target.Host)
	if err != nil {
		return 0, nil, err
	}
	defer conn.Close()
	if t.timeout > 0 {
		_ = conn.SetDeadline(time.Now().Add(t.timeout))
	}

	if err := writeJSONRequest(conn, target, body, headers); err != nil {
		return 0, nil, err
	}

	resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
	if err != nil {
		return 0, nil, err
	}
	defer resp.Body.Close()

	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		return resp.StatusCode, nil, fmt.Errorf("read response body: %w", err)
	}
	return resp.StatusCode, raw, nil
}

func writeJSONRequest(w io.Writer, target *url.URL, body []byte, headers map[string]string) error {
	path := target.RequestURI()
	if path == "" {
		path = "/"
	}

	var buf bytes.Buffer
	fmt.Fprintf(&buf, "POST %s HTTP/1.1\r\n", path)
	writeHeader(&buf, "Host", target.Host)
	writeHeader(&buf, "User-Agent", headerValue(headers, "User-Agent", "python-requests/2.32.4"))
	writeHeader(&buf, "Accept-Encoding", headerValue(headers, "Accept-Encoding", "identity"))
	writeHeader(&buf, "Accept", headerValue(headers, "Accept", "*/*"))
	writeHeader(&buf, "Connection", headerValue(headers, "Connection", "keep-alive"))
	writeHeader(&buf, "Content-Length", fmt.Sprintf("%d", len(body)))
	writeHeader(&buf, "Content-Type", headerValue(headers, "Content-Type", "application/json"))

	written := map[string]bool{
		"host": true, "user-agent": true, "accept-encoding": true, "accept": true,
		"connection": true, "content-length": true, "content-type": true,
	}
	for _, name := range []string{"orgCode", "deviceSN", "ts", "userName", "password"} {
		if value, ok := headers[name]; ok {
			writeHeader(&buf, name, value)
			written[strings.ToLower(name)] = true
		}
	}

	var rest []string
	for name := range headers {
		if !written[strings.ToLower(name)] {
			rest = append(rest, name)
		}
	}
	sort.Strings(rest)
	for _, name := range rest {
		writeHeader(&buf, name, headers[name])
	}
	buf.WriteString("\r\n")
	buf.Write(body)

	_, err := w.Write(buf.Bytes())
	return err
}

func writeHeader(buf *bytes.Buffer, name, value string) {
	fmt.Fprintf(buf, "%s: %s\r\n", name, value)
}

func headerValue(headers map[string]string, name, fallback string) string {
	if value, ok := headers[name]; ok {
		return value
	}
	return fallback
}

func dialWithContext(ctx context.Context, dialer proxy.Dialer, network, address string) (net.Conn, error) {
	if contextDialer, ok := dialer.(proxy.ContextDialer); ok {
		return contextDialer.DialContext(ctx, network, address)
	}
	type result struct {
		conn net.Conn
		err  error
	}
	ch := make(chan result, 1)
	go func() {
		conn, err := dialer.Dial(network, address)
		ch <- result{conn: conn, err: err}
	}()
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case result := <-ch:
		return result.conn, result.err
	}
}

func socks5Dialer(rawProxy string) (proxy.Dialer, error) {
	proxyURL, err := normalizeSocks5Proxy(rawProxy)
	if err != nil {
		return nil, err
	}

	dialer, err := proxy.SOCKS5("tcp", proxyURL.Host, nil, proxy.Direct)
	if err != nil {
		return nil, fmt.Errorf("create socks5 dialer: %w", err)
	}
	return dialer, nil
}

func normalizeSocks5Proxy(rawProxy string) (*url.URL, error) {
	rawProxy = strings.TrimSpace(rawProxy)
	if !strings.Contains(rawProxy, "://") {
		rawProxy = "socks5://" + rawProxy
	}

	proxyURL, err := url.Parse(rawProxy)
	if err != nil {
		return nil, fmt.Errorf("parse socks5 proxy: %w", err)
	}
	if proxyURL.Scheme != "socks5" && proxyURL.Scheme != "socks5h" {
		return nil, fmt.Errorf("unsupported proxy scheme %q", proxyURL.Scheme)
	}
	if proxyURL.User != nil {
		return nil, fmt.Errorf("socks5 proxy authentication is not supported yet")
	}
	if _, _, err := net.SplitHostPort(proxyURL.Host); err != nil {
		return nil, fmt.Errorf("invalid socks5 proxy host %q: %w", proxyURL.Host, err)
	}
	return proxyURL, nil
}