From 18ba5e6a308897bbfd4e40571cfb5fd5305faa83 Mon Sep 17 00:00:00 2001
From: Tnze <cjd001113@outlook.com>
Date: Wed, 15 Feb 2023 19:32:59 +0800
Subject: [PATCH] proxy using support. close #231

---
 bot/mcbot.go | 29 +++++++++++++++--------------
 net/conn.go  | 27 ++++++++++++++++++++++-----
 2 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/bot/mcbot.go b/bot/mcbot.go
index 7eb247a..86a2a59 100644
--- a/bot/mcbot.go
+++ b/bot/mcbot.go
@@ -10,12 +10,13 @@ import (
 	"net"
 	"strconv"
 
+	"github.com/google/uuid"
+
 	"github.com/Tnze/go-mc/chat"
 	"github.com/Tnze/go-mc/data/packetid"
 	mcnet "github.com/Tnze/go-mc/net"
 	pk "github.com/Tnze/go-mc/net/packet"
 	"github.com/Tnze/go-mc/yggdrasil/user"
-	"github.com/google/uuid"
 )
 
 // ProtocolVersion is the protocol version number of minecraft net protocol
@@ -25,8 +26,8 @@ const (
 )
 
 type JoinOptions struct {
-	Dialer  *net.Dialer
-	Context context.Context
+	MCDialer mcnet.MCDialer
+	Context  context.Context
 
 	// Indicate not to fetch and sending player's PubKey
 	NoPublicKey bool
@@ -40,22 +41,22 @@ type JoinOptions struct {
 // Using roughly the same way to parse address as minecraft.
 func (c *Client) JoinServer(addr string) (err error) {
 	return c.join(addr, JoinOptions{
-		Context: context.Background(),
-		Dialer:  (*net.Dialer)(&mcnet.DefaultDialer),
+		Context:  context.Background(),
+		MCDialer: &mcnet.DefaultDialer,
 	})
 }
 
-// JoinServerWithDialer is similar to JoinServer but using a Dialer.
+// JoinServerWithDialer is similar to JoinServer but using a net.Dialer.
 func (c *Client) JoinServerWithDialer(dialer *net.Dialer, addr string) (err error) {
 	return c.join(addr, JoinOptions{
-		Context: context.Background(),
-		Dialer:  dialer,
+		Context:  context.Background(),
+		MCDialer: (*mcnet.Dialer)(dialer),
 	})
 }
 
 func (c *Client) JoinServerWithOptions(addr string, options JoinOptions) (err error) {
-	if options.Dialer == nil {
-		options.Dialer = (*net.Dialer)(&mcnet.DefaultDialer)
+	if options.MCDialer == nil {
+		options.MCDialer = &mcnet.DefaultDialer
 	}
 	if options.Context == nil {
 		options.Context = context.Background()
@@ -65,7 +66,9 @@ func (c *Client) JoinServerWithOptions(addr string, options JoinOptions) (err er
 
 func (c *Client) join(addr string, options JoinOptions) error {
 	const Handshake = 0x00
-	// Split Host and Port
+
+	// Split Host and Port. The DialMCContext will do this once,
+	// but we need the result for sending handshake packet here.
 	host, portStr, err := net.SplitHostPort(addr)
 	var port uint64
 	if err != nil {
@@ -85,9 +88,7 @@ func (c *Client) join(addr string, options JoinOptions) error {
 	}
 
 	// Dial connection
-	d := (*mcnet.Dialer)(options.Dialer)
-	ctx := options.Context
-	c.Conn, err = d.DialMCContext(ctx, addr)
+	c.Conn, err = options.MCDialer.DialMCContext(options.Context, addr)
 	if err != nil {
 		return LoginErr{"connect server", err}
 	}
diff --git a/net/conn.go b/net/conn.go
index 037f991..db87e75 100644
--- a/net/conn.go
+++ b/net/conn.go
@@ -62,6 +62,23 @@ func DialMCTimeout(addr string, timeout time.Duration) (*Conn, error) {
 	return DefaultDialer.DialMCContext(ctx, addr)
 }
 
+// MCDialer provide DialMCContext method, can be used to dial a minecraft server.
+// [Dialer] is its default implementation, and support SRV lookup.
+//
+// Typically, if you want to use built-in proxies or custom dialer,
+// you can hook go-mc/bot package by implement this interface.
+// When implementing a custom MCDialer, SRV lookup is optional.
+type MCDialer interface {
+	// The DialMCContext dial TCP connection to a minecraft server, and warp the net.Conn by calling [WrapConn].
+	DialMCContext(ctx context.Context, addr string) (*Conn, error)
+}
+
+// Dialer implements MCDialer interface.
+//
+// It can be easily convert from net.Dialer.
+//
+//	dialer := net.Dialer{}
+//	mcDialer := (*Dialer)(&dialer)
 type Dialer net.Dialer
 
 func (d *Dialer) resolver() *net.Resolver {
@@ -82,23 +99,23 @@ func (d *Dialer) DialMCContext(ctx context.Context, addr string) (*Conn, error)
 			return nil, err
 		}
 	}
-	var ras []string
+	var addresses []string
 	if port == "" {
 		// We look up SRV only if the port is not specified
 		_, srvRecords, err := d.resolver().LookupSRV(ctx, "minecraft", "tcp", host)
 		if err == nil {
 			for _, record := range srvRecords {
 				addr := net.JoinHostPort(record.Target, strconv.Itoa(int(record.Port)))
-				ras = append(ras, addr)
+				addresses = append(addresses, addr)
 			}
 		}
 		// Whatever the SRV records is found,
 		addr = net.JoinHostPort(addr, strconv.Itoa(DefaultPort))
 	}
-	ras = append(ras, addr)
+	addresses = append(addresses, addr)
 
 	var firstErr error
-	for i, addr := range ras {
+	for i, addr := range addresses {
 		select {
 		case <-ctx.Done():
 			return nil, context.Canceled
@@ -106,7 +123,7 @@ func (d *Dialer) DialMCContext(ctx context.Context, addr string) (*Conn, error)
 		}
 		dialCtx := ctx
 		if deadline, hasDeadline := ctx.Deadline(); hasDeadline {
-			partialDeadline, err := partialDeadline(time.Now(), deadline, len(ras)-i)
+			partialDeadline, err := partialDeadline(time.Now(), deadline, len(addresses)-i)
 			if err != nil {
 				// Ran out of time.
 				if firstErr == nil {