chore: 准备发布 v1.0.0 版本
This commit is contained in:
@@ -0,0 +1,615 @@
|
||||
package geoip
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/oschwald/geoip2-golang"
|
||||
)
|
||||
|
||||
// Region codes
|
||||
const (
|
||||
RegionJP = "jp"
|
||||
RegionKR = "kr"
|
||||
RegionUS = "us"
|
||||
RegionHK = "hk"
|
||||
RegionTW = "tw"
|
||||
RegionOther = "other"
|
||||
)
|
||||
|
||||
// Default GeoIP database download URL
|
||||
const (
|
||||
DefaultGeoIPURL = "https://github.com/P3TERX/GeoLite.mmdb/raw/download/GeoLite2-Country.mmdb"
|
||||
)
|
||||
|
||||
// RegionInfo contains region details
|
||||
type RegionInfo struct {
|
||||
Code string // "jp", "kr", "us", "hk", "tw", "other"
|
||||
Country string // Full country name
|
||||
ISOCode string // ISO country code
|
||||
}
|
||||
|
||||
// Lookup provides GeoIP lookup functionality
|
||||
type Lookup struct {
|
||||
db *geoip2.Reader
|
||||
mu sync.RWMutex
|
||||
path string
|
||||
updateInterval time.Duration
|
||||
stopChan chan struct{}
|
||||
updateOnce sync.Once
|
||||
}
|
||||
|
||||
// EnsureDatabase checks if the GeoIP database exists, and downloads it if not
|
||||
func EnsureDatabase(dbPath string) error {
|
||||
if dbPath == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if file already exists and is valid
|
||||
info, err := os.Stat(dbPath)
|
||||
if err == nil {
|
||||
if !info.Mode().IsRegular() {
|
||||
return fmt.Errorf("geoip database path is not a file: %s", dbPath)
|
||||
}
|
||||
if info.Size() > 0 {
|
||||
return nil // File exists and has content
|
||||
}
|
||||
} else if !os.IsNotExist(err) {
|
||||
return fmt.Errorf("stat geoip database: %w", err)
|
||||
}
|
||||
|
||||
log.Printf("📥 GeoIP database not found at %s, downloading...", dbPath)
|
||||
|
||||
// Create parent directory if needed
|
||||
dir := filepath.Dir(dbPath)
|
||||
if dir != "." {
|
||||
if err := os.MkdirAll(dir, 0755); err != nil {
|
||||
return fmt.Errorf("create directory: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Download with timeout
|
||||
client := &http.Client{Timeout: 60 * time.Second}
|
||||
req, err := http.NewRequest(http.MethodGet, DefaultGeoIPURL, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create download request: %w", err)
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("download failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("download failed: unexpected status %s", resp.Status)
|
||||
}
|
||||
|
||||
// Download to temporary file
|
||||
tempFile, err := os.CreateTemp(dir, ".geoip-*.mmdb")
|
||||
if err != nil {
|
||||
return fmt.Errorf("create temp file: %w", err)
|
||||
}
|
||||
tempPath := tempFile.Name()
|
||||
cleanup := true
|
||||
defer func() {
|
||||
if tempFile != nil {
|
||||
tempFile.Close()
|
||||
}
|
||||
if cleanup {
|
||||
os.Remove(tempPath)
|
||||
}
|
||||
}()
|
||||
|
||||
// Copy with progress tracking
|
||||
progress := &progressWriter{total: resp.ContentLength}
|
||||
reader := io.TeeReader(resp.Body, progress)
|
||||
written, err := io.Copy(tempFile, reader)
|
||||
if err != nil {
|
||||
return fmt.Errorf("download failed: %w", err)
|
||||
}
|
||||
|
||||
// Verify download completeness
|
||||
if resp.ContentLength > 0 && written < resp.ContentLength {
|
||||
return fmt.Errorf("incomplete download (%d/%d bytes)", written, resp.ContentLength)
|
||||
}
|
||||
|
||||
// Sync and close temp file
|
||||
if err := tempFile.Sync(); err != nil {
|
||||
return fmt.Errorf("sync temp file: %w", err)
|
||||
}
|
||||
if err := tempFile.Close(); err != nil {
|
||||
return fmt.Errorf("close temp file: %w", err)
|
||||
}
|
||||
tempFile = nil
|
||||
|
||||
// Validate MMDB format
|
||||
if err := validateMMDB(tempPath); err != nil {
|
||||
return fmt.Errorf("validation failed: %w", err)
|
||||
}
|
||||
|
||||
// Atomic rename
|
||||
if err := os.Rename(tempPath, dbPath); err != nil {
|
||||
return fmt.Errorf("rename failed: %w", err)
|
||||
}
|
||||
cleanup = false
|
||||
|
||||
log.Printf("✅ GeoIP database downloaded successfully to %s", dbPath)
|
||||
return nil
|
||||
}
|
||||
|
||||
// progressWriter tracks download progress
|
||||
type progressWriter struct {
|
||||
total int64
|
||||
downloaded int64
|
||||
lastPercent int64
|
||||
lastLog time.Time
|
||||
}
|
||||
|
||||
func (p *progressWriter) Write(b []byte) (int, error) {
|
||||
n := len(b)
|
||||
p.downloaded += int64(n)
|
||||
|
||||
now := time.Now()
|
||||
if p.total > 0 {
|
||||
percent := p.downloaded * 100 / p.total
|
||||
if percent >= 100 || percent >= p.lastPercent+10 || now.Sub(p.lastLog) >= 3*time.Second {
|
||||
log.Printf(" Progress: %d%% (%d/%d bytes)", percent, p.downloaded, p.total)
|
||||
p.lastPercent = percent
|
||||
p.lastLog = now
|
||||
}
|
||||
} else if now.Sub(p.lastLog) >= 3*time.Second {
|
||||
log.Printf(" Downloaded: %d bytes", p.downloaded)
|
||||
p.lastLog = now
|
||||
}
|
||||
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// validateMMDB performs basic validation of MMDB file format
|
||||
func validateMMDB(path string) error {
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
info, err := file.Stat()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if info.Size() < 1024 {
|
||||
return fmt.Errorf("file too small (%d bytes)", info.Size())
|
||||
}
|
||||
|
||||
// Check for MaxMind metadata in the last 8KB
|
||||
const tailSize int64 = 8192
|
||||
readSize := tailSize
|
||||
if info.Size() < readSize {
|
||||
readSize = info.Size()
|
||||
}
|
||||
if _, err := file.Seek(-readSize, io.SeekEnd); err != nil {
|
||||
return err
|
||||
}
|
||||
buf := make([]byte, readSize)
|
||||
if _, err := io.ReadFull(file, buf); err != nil && err != io.ErrUnexpectedEOF {
|
||||
return err
|
||||
}
|
||||
if !bytes.Contains(buf, []byte("MaxMind.com")) {
|
||||
return fmt.Errorf("missing MaxMind metadata")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// downloadFile downloads a file from URL to the specified path
|
||||
func downloadFile(filepath string, url string) error {
|
||||
// Create the file
|
||||
out, err := os.Create(filepath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer out.Close()
|
||||
|
||||
// Get the data
|
||||
resp, err := http.Get(url)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Check server response
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("bad status: %s", resp.Status)
|
||||
}
|
||||
|
||||
// Writer the body to file
|
||||
size, err := io.Copy(out, resp.Body)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
log.Printf(" Downloaded %d bytes", size)
|
||||
return nil
|
||||
}
|
||||
|
||||
// New creates a new GeoIP lookup instance
|
||||
func New(dbPath string) (*Lookup, error) {
|
||||
return NewWithAutoUpdate(dbPath, 0)
|
||||
}
|
||||
|
||||
// NewWithAutoUpdate creates a new GeoIP lookup instance with auto-update support
|
||||
func NewWithAutoUpdate(dbPath string, updateInterval time.Duration) (*Lookup, error) {
|
||||
if dbPath == "" {
|
||||
return &Lookup{}, nil
|
||||
}
|
||||
|
||||
// Ensure database exists (download if needed)
|
||||
if err := EnsureDatabase(dbPath); err != nil {
|
||||
return nil, fmt.Errorf("ensure database: %w", err)
|
||||
}
|
||||
|
||||
db, err := geoip2.Open(dbPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
lookup := &Lookup{
|
||||
db: db,
|
||||
path: dbPath,
|
||||
updateInterval: updateInterval,
|
||||
stopChan: make(chan struct{}),
|
||||
}
|
||||
|
||||
// Start auto-update goroutine if interval is set
|
||||
if updateInterval > 0 {
|
||||
go lookup.autoUpdateLoop()
|
||||
log.Printf("🔄 GeoIP auto-update enabled (interval: %v)", updateInterval)
|
||||
}
|
||||
|
||||
return lookup, nil
|
||||
}
|
||||
|
||||
// autoUpdateLoop periodically updates the GeoIP database
|
||||
func (l *Lookup) autoUpdateLoop() {
|
||||
ticker := time.NewTicker(l.updateInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
if err := l.Update(); err != nil {
|
||||
log.Printf("⚠️ GeoIP auto-update failed: %v", err)
|
||||
}
|
||||
case <-l.stopChan:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Update downloads and reloads the GeoIP database
|
||||
func (l *Lookup) Update() error {
|
||||
log.Printf("🔄 Updating GeoIP database...")
|
||||
|
||||
// Download to temporary file
|
||||
tempPath := l.path + ".update"
|
||||
if err := downloadDatabase(tempPath); err != nil {
|
||||
return fmt.Errorf("download failed: %w", err)
|
||||
}
|
||||
defer os.Remove(tempPath) // Clean up temp file
|
||||
|
||||
// Validate the downloaded database
|
||||
if err := validateMMDB(tempPath); err != nil {
|
||||
return fmt.Errorf("validation failed: %w", err)
|
||||
}
|
||||
|
||||
// Open new database
|
||||
newDB, err := geoip2.Open(tempPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("open new database: %w", err)
|
||||
}
|
||||
|
||||
// Hot-swap the database
|
||||
l.mu.Lock()
|
||||
oldDB := l.db
|
||||
l.db = newDB
|
||||
l.mu.Unlock()
|
||||
|
||||
// Close old database
|
||||
if oldDB != nil {
|
||||
oldDB.Close()
|
||||
}
|
||||
|
||||
// Replace the old file with new one
|
||||
if err := os.Rename(tempPath, l.path); err != nil {
|
||||
log.Printf("⚠️ Failed to replace database file: %v (using in-memory version)", err)
|
||||
}
|
||||
|
||||
log.Printf("✅ GeoIP database updated successfully")
|
||||
return nil
|
||||
}
|
||||
|
||||
// downloadDatabase downloads the GeoIP database to the specified path
|
||||
func downloadDatabase(filepath string) error {
|
||||
// Create parent directory if needed
|
||||
dir := filepath[:strings.LastIndex(filepath, "/")]
|
||||
if dir != "" && dir != "." {
|
||||
if err := os.MkdirAll(dir, 0755); err != nil {
|
||||
return fmt.Errorf("create directory: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Download with timeout
|
||||
client := &http.Client{Timeout: 60 * time.Second}
|
||||
req, err := http.NewRequest(http.MethodGet, DefaultGeoIPURL, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create download request: %w", err)
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("unexpected status %s", resp.Status)
|
||||
}
|
||||
|
||||
// Create temp file
|
||||
tempFile, err := os.CreateTemp(dir, ".geoip-download-*.mmdb")
|
||||
if err != nil {
|
||||
return fmt.Errorf("create temp file: %w", err)
|
||||
}
|
||||
tempPath := tempFile.Name()
|
||||
cleanup := true
|
||||
defer func() {
|
||||
if tempFile != nil {
|
||||
tempFile.Close()
|
||||
}
|
||||
if cleanup {
|
||||
os.Remove(tempPath)
|
||||
}
|
||||
}()
|
||||
|
||||
// Copy with progress tracking
|
||||
progress := &progressWriter{total: resp.ContentLength}
|
||||
reader := io.TeeReader(resp.Body, progress)
|
||||
written, err := io.Copy(tempFile, reader)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Verify download completeness
|
||||
if resp.ContentLength > 0 && written < resp.ContentLength {
|
||||
return fmt.Errorf("incomplete download (%d/%d bytes)", written, resp.ContentLength)
|
||||
}
|
||||
|
||||
// Sync and close
|
||||
if err := tempFile.Sync(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tempFile.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
tempFile = nil
|
||||
|
||||
// Rename to target path
|
||||
if err := os.Rename(tempPath, filepath); err != nil {
|
||||
return err
|
||||
}
|
||||
cleanup = false
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close closes the GeoIP database and stops auto-update
|
||||
func (l *Lookup) Close() error {
|
||||
// Stop auto-update goroutine
|
||||
l.updateOnce.Do(func() {
|
||||
if l.stopChan != nil {
|
||||
close(l.stopChan)
|
||||
}
|
||||
})
|
||||
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
if l.db != nil {
|
||||
return l.db.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// IsEnabled returns true if GeoIP lookup is available
|
||||
func (l *Lookup) IsEnabled() bool {
|
||||
l.mu.RLock()
|
||||
defer l.mu.RUnlock()
|
||||
return l.db != nil
|
||||
}
|
||||
|
||||
// LookupIP returns region info for an IP address
|
||||
func (l *Lookup) LookupIP(ipStr string) RegionInfo {
|
||||
l.mu.RLock()
|
||||
defer l.mu.RUnlock()
|
||||
|
||||
if l.db == nil {
|
||||
return RegionInfo{Code: RegionOther, Country: "Unknown", ISOCode: ""}
|
||||
}
|
||||
|
||||
ip := net.ParseIP(ipStr)
|
||||
if ip == nil {
|
||||
return RegionInfo{Code: RegionOther, Country: "Unknown", ISOCode: ""}
|
||||
}
|
||||
|
||||
record, err := l.db.Country(ip)
|
||||
if err != nil {
|
||||
return RegionInfo{Code: RegionOther, Country: "Unknown", ISOCode: ""}
|
||||
}
|
||||
|
||||
isoCode := record.Country.IsoCode
|
||||
country := record.Country.Names["en"]
|
||||
if country == "" {
|
||||
country = isoCode
|
||||
}
|
||||
|
||||
return RegionInfo{
|
||||
Code: isoCodeToRegion(isoCode),
|
||||
Country: country,
|
||||
ISOCode: isoCode,
|
||||
}
|
||||
}
|
||||
|
||||
// LookupURI extracts server from URI and returns region info
|
||||
func (l *Lookup) LookupURI(uri string) RegionInfo {
|
||||
host := extractHostFromURI(uri)
|
||||
if host == "" {
|
||||
return RegionInfo{Code: RegionOther, Country: "Unknown", ISOCode: ""}
|
||||
}
|
||||
|
||||
// Resolve hostname to IP if needed
|
||||
ip := net.ParseIP(host)
|
||||
if ip == nil {
|
||||
// It's a hostname, try to resolve
|
||||
ips, err := net.LookupIP(host)
|
||||
if err != nil || len(ips) == 0 {
|
||||
return RegionInfo{Code: RegionOther, Country: "Unknown", ISOCode: ""}
|
||||
}
|
||||
host = ips[0].String()
|
||||
}
|
||||
|
||||
return l.LookupIP(host)
|
||||
}
|
||||
|
||||
// extractHostFromURI extracts the host/IP from various proxy URI formats
|
||||
func extractHostFromURI(uri string) string {
|
||||
// Handle different URI schemes
|
||||
lowerURI := strings.ToLower(uri)
|
||||
|
||||
// vmess:// vless:// trojan:// ss:// ssr:// hysteria:// hysteria2:// hy2://
|
||||
if strings.HasPrefix(lowerURI, "vmess://") ||
|
||||
strings.HasPrefix(lowerURI, "vless://") ||
|
||||
strings.HasPrefix(lowerURI, "trojan://") ||
|
||||
strings.HasPrefix(lowerURI, "hysteria://") ||
|
||||
strings.HasPrefix(lowerURI, "hysteria2://") ||
|
||||
strings.HasPrefix(lowerURI, "hy2://") {
|
||||
// Standard URL format: scheme://user@host:port?params#fragment
|
||||
parsed, err := url.Parse(uri)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return parsed.Hostname()
|
||||
}
|
||||
|
||||
if strings.HasPrefix(lowerURI, "ss://") {
|
||||
// Shadowsocks format: ss://base64(method:password)@host:port#name
|
||||
// or ss://base64@host:port#name
|
||||
return extractSSHost(uri)
|
||||
}
|
||||
|
||||
if strings.HasPrefix(lowerURI, "ssr://") {
|
||||
// SSR format is base64 encoded
|
||||
return extractSSRHost(uri)
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
func extractSSHost(uri string) string {
|
||||
// Remove ss:// prefix
|
||||
content := strings.TrimPrefix(uri, "ss://")
|
||||
|
||||
// Remove fragment (#name)
|
||||
if idx := strings.Index(content, "#"); idx != -1 {
|
||||
content = content[:idx]
|
||||
}
|
||||
|
||||
// Check if it's the new format: base64@host:port
|
||||
if atIdx := strings.LastIndex(content, "@"); atIdx != -1 {
|
||||
hostPort := content[atIdx+1:]
|
||||
if colonIdx := strings.LastIndex(hostPort, ":"); colonIdx != -1 {
|
||||
return hostPort[:colonIdx]
|
||||
}
|
||||
return hostPort
|
||||
}
|
||||
|
||||
// Old format: entire content is base64
|
||||
return ""
|
||||
}
|
||||
|
||||
func extractSSRHost(uri string) string {
|
||||
// SSR is complex, skip for now - will be marked as "other"
|
||||
return ""
|
||||
}
|
||||
|
||||
// isoCodeToRegion maps ISO country codes to our region codes
|
||||
func isoCodeToRegion(isoCode string) string {
|
||||
switch strings.ToUpper(isoCode) {
|
||||
case "JP":
|
||||
return RegionJP
|
||||
case "KR":
|
||||
return RegionKR
|
||||
case "US":
|
||||
return RegionUS
|
||||
case "HK":
|
||||
return RegionHK
|
||||
case "TW":
|
||||
return RegionTW
|
||||
default:
|
||||
return RegionOther
|
||||
}
|
||||
}
|
||||
|
||||
// AllRegions returns all supported region codes
|
||||
func AllRegions() []string {
|
||||
return []string{RegionJP, RegionKR, RegionUS, RegionHK, RegionTW, RegionOther}
|
||||
}
|
||||
|
||||
// RegionName returns the display name for a region code
|
||||
func RegionName(code string) string {
|
||||
switch code {
|
||||
case RegionJP:
|
||||
return "Japan"
|
||||
case RegionKR:
|
||||
return "Korea"
|
||||
case RegionUS:
|
||||
return "USA"
|
||||
case RegionHK:
|
||||
return "Hong Kong"
|
||||
case RegionTW:
|
||||
return "Taiwan"
|
||||
case RegionOther:
|
||||
return "Other"
|
||||
default:
|
||||
return "Unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// RegionEmoji returns the flag emoji for a region code
|
||||
func RegionEmoji(code string) string {
|
||||
switch code {
|
||||
case RegionJP:
|
||||
return "🇯🇵"
|
||||
case RegionKR:
|
||||
return "🇰🇷"
|
||||
case RegionUS:
|
||||
return "🇺🇸"
|
||||
case RegionHK:
|
||||
return "🇭🇰"
|
||||
case RegionTW:
|
||||
return "🇹🇼"
|
||||
case RegionOther:
|
||||
return "🌍"
|
||||
default:
|
||||
return "❓"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,271 @@
|
||||
package geoip
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// RouterConfig holds configuration for the GeoIP router
|
||||
type RouterConfig struct {
|
||||
Listen string
|
||||
Port uint16
|
||||
Username string
|
||||
Password string
|
||||
}
|
||||
|
||||
// PoolDialer is an interface for dialing through a specific pool
|
||||
type PoolDialer interface {
|
||||
DialContext(ctx context.Context, network, address string) (net.Conn, error)
|
||||
}
|
||||
|
||||
// Router handles HTTP proxy requests with path-based region routing
|
||||
type Router struct {
|
||||
cfg RouterConfig
|
||||
pools map[string]PoolDialer // region -> dialer
|
||||
global PoolDialer // default pool for requests without region path
|
||||
server *http.Server
|
||||
mu sync.RWMutex
|
||||
logger *log.Logger
|
||||
}
|
||||
|
||||
// NewRouter creates a new GeoIP router
|
||||
func NewRouter(cfg RouterConfig, logger *log.Logger) *Router {
|
||||
if logger == nil {
|
||||
logger = log.Default()
|
||||
}
|
||||
return &Router{
|
||||
cfg: cfg,
|
||||
pools: make(map[string]PoolDialer),
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
// SetPool registers a pool dialer for a specific region
|
||||
func (r *Router) SetPool(region string, dialer PoolDialer) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.pools[region] = dialer
|
||||
}
|
||||
|
||||
// SetGlobalPool sets the default pool for requests without region path
|
||||
func (r *Router) SetGlobalPool(dialer PoolDialer) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.global = dialer
|
||||
}
|
||||
|
||||
// Start starts the GeoIP router HTTP server
|
||||
func (r *Router) Start(ctx context.Context) error {
|
||||
addr := fmt.Sprintf("%s:%d", r.cfg.Listen, r.cfg.Port)
|
||||
|
||||
r.server = &http.Server{
|
||||
Addr: addr,
|
||||
Handler: r,
|
||||
}
|
||||
|
||||
go func() {
|
||||
r.logger.Printf("🌐 GeoIP Router started on %s", addr)
|
||||
r.logger.Println(" Routes: /jp, /kr, /us, /hk, /tw, /other (default: all nodes)")
|
||||
if err := r.server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
r.logger.Printf("GeoIP router error: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
go func() {
|
||||
<-ctx.Done()
|
||||
r.Stop()
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop stops the GeoIP router
|
||||
func (r *Router) Stop() error {
|
||||
if r.server != nil {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
return r.server.Shutdown(ctx)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ServeHTTP handles incoming HTTP proxy requests
|
||||
func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
|
||||
// Check authentication if configured
|
||||
if r.cfg.Username != "" {
|
||||
user, pass, ok := req.BasicAuth()
|
||||
if !ok || user != r.cfg.Username || pass != r.cfg.Password {
|
||||
w.Header().Set("Proxy-Authenticate", `Basic realm="Proxy"`)
|
||||
http.Error(w, "Proxy authentication required", http.StatusProxyAuthRequired)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Extract region from path
|
||||
region, targetHost := r.parseRequest(req)
|
||||
|
||||
// Get the appropriate pool
|
||||
r.mu.RLock()
|
||||
var dialer PoolDialer
|
||||
if region != "" {
|
||||
dialer = r.pools[region]
|
||||
}
|
||||
if dialer == nil {
|
||||
dialer = r.global
|
||||
}
|
||||
r.mu.RUnlock()
|
||||
|
||||
if dialer == nil {
|
||||
http.Error(w, "No proxy pool available", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
|
||||
if req.Method == http.MethodConnect {
|
||||
r.handleConnect(w, req, dialer, targetHost)
|
||||
} else {
|
||||
r.handleHTTP(w, req, dialer, targetHost)
|
||||
}
|
||||
}
|
||||
|
||||
// parseRequest extracts region and target host from the request
|
||||
func (r *Router) parseRequest(req *http.Request) (region, targetHost string) {
|
||||
// For CONNECT requests, the host is in req.Host
|
||||
// For regular requests, check the path prefix
|
||||
|
||||
if req.Method == http.MethodConnect {
|
||||
// CONNECT requests: check if host starts with region prefix
|
||||
// e.g., CONNECT jp/example.com:443 or just example.com:443
|
||||
host := req.Host
|
||||
for _, reg := range AllRegions() {
|
||||
prefix := reg + "/"
|
||||
if strings.HasPrefix(host, prefix) {
|
||||
return reg, strings.TrimPrefix(host, prefix)
|
||||
}
|
||||
}
|
||||
return "", host
|
||||
}
|
||||
|
||||
// For regular HTTP requests, check URL path
|
||||
path := req.URL.Path
|
||||
for _, reg := range AllRegions() {
|
||||
prefix := "/" + reg + "/"
|
||||
if strings.HasPrefix(path, prefix) {
|
||||
// Rewrite the path
|
||||
req.URL.Path = "/" + strings.TrimPrefix(path, prefix)
|
||||
return reg, req.Host
|
||||
}
|
||||
// Also check for exact match like /jp
|
||||
if path == "/"+reg {
|
||||
req.URL.Path = "/"
|
||||
return reg, req.Host
|
||||
}
|
||||
}
|
||||
|
||||
return "", req.Host
|
||||
}
|
||||
|
||||
// handleConnect handles HTTPS CONNECT tunneling
|
||||
func (r *Router) handleConnect(w http.ResponseWriter, req *http.Request, dialer PoolDialer, targetHost string) {
|
||||
ctx, cancel := context.WithTimeout(req.Context(), 30*time.Second)
|
||||
defer cancel()
|
||||
|
||||
targetConn, err := dialer.DialContext(ctx, "tcp", targetHost)
|
||||
if err != nil {
|
||||
http.Error(w, fmt.Sprintf("Failed to connect: %v", err), http.StatusBadGateway)
|
||||
return
|
||||
}
|
||||
defer targetConn.Close()
|
||||
|
||||
hijacker, ok := w.(http.Hijacker)
|
||||
if !ok {
|
||||
http.Error(w, "Hijacking not supported", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
clientConn, _, err := hijacker.Hijack()
|
||||
if err != nil {
|
||||
http.Error(w, fmt.Sprintf("Hijack failed: %v", err), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
defer clientConn.Close()
|
||||
|
||||
// Send 200 Connection Established
|
||||
clientConn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))
|
||||
|
||||
// Bidirectional copy
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
io.Copy(targetConn, clientConn)
|
||||
}()
|
||||
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
io.Copy(clientConn, targetConn)
|
||||
}()
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// handleHTTP handles regular HTTP requests
|
||||
func (r *Router) handleHTTP(w http.ResponseWriter, req *http.Request, dialer PoolDialer, targetHost string) {
|
||||
ctx, cancel := context.WithTimeout(req.Context(), 30*time.Second)
|
||||
defer cancel()
|
||||
|
||||
// Create a new request to the target
|
||||
targetURL := req.URL
|
||||
if targetURL.Host == "" {
|
||||
targetURL.Host = targetHost
|
||||
}
|
||||
if targetURL.Scheme == "" {
|
||||
targetURL.Scheme = "http"
|
||||
}
|
||||
|
||||
outReq, err := http.NewRequestWithContext(ctx, req.Method, targetURL.String(), req.Body)
|
||||
if err != nil {
|
||||
http.Error(w, fmt.Sprintf("Failed to create request: %v", err), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
// Copy headers
|
||||
for key, values := range req.Header {
|
||||
for _, value := range values {
|
||||
outReq.Header.Add(key, value)
|
||||
}
|
||||
}
|
||||
|
||||
// Remove hop-by-hop headers
|
||||
outReq.Header.Del("Proxy-Connection")
|
||||
outReq.Header.Del("Proxy-Authorization")
|
||||
|
||||
// Create transport with custom dialer
|
||||
transport := &http.Transport{
|
||||
DialContext: dialer.DialContext,
|
||||
}
|
||||
|
||||
resp, err := transport.RoundTrip(outReq)
|
||||
if err != nil {
|
||||
http.Error(w, fmt.Sprintf("Request failed: %v", err), http.StatusBadGateway)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Copy response headers
|
||||
for key, values := range resp.Header {
|
||||
for _, value := range values {
|
||||
w.Header().Add(key, value)
|
||||
}
|
||||
}
|
||||
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
io.Copy(w, resp.Body)
|
||||
}
|
||||
Reference in New Issue
Block a user