From 177ecb3de2f9ca8e23c3eeb19b6875a0c6593355 Mon Sep 17 00:00:00 2001
From: Keuin <keuinx@gmail.com>
Date: Thu, 8 Sep 2022 13:27:37 +0800
Subject: Support specify IP network type (ipv4, ipv6, in arbitrary combination
 and priority)

---
 bilibili/client.go    | 28 ++++++++++++++++++++++++----
 bilibili/netprobe.go  | 41 +++++++++++++++++++++++++++++++++++++++++
 bilibili/request.go   | 26 +++++++++++++++++++++++++-
 bilibili/streaming.go |  2 +-
 4 files changed, 91 insertions(+), 6 deletions(-)
 create mode 100644 bilibili/netprobe.go

(limited to 'bilibili')

diff --git a/bilibili/client.go b/bilibili/client.go
index fcea395..8097115 100644
--- a/bilibili/client.go
+++ b/bilibili/client.go
@@ -7,6 +7,7 @@ package bilibili
 import (
 	"context"
 	"log"
+	"net"
 	"net/http"
 	"os"
 )
@@ -21,25 +22,44 @@ type Bilibili struct {
 	userAgent string
 	http      *http.Client
 	loggerCommon
-	ctx context.Context
+	ctx      context.Context
+	netTypes []IpNetType
 }
 
-func NewBilibiliWithContext(ctx context.Context) Bilibili {
+func NewBilibiliWithContext(ctx context.Context, netTypes []IpNetType) Bilibili {
 	logger := loggerCommon{
 		debug: log.New(os.Stderr, "DEBUG: ", log.Ldate|log.Ltime|log.Lshortfile),
 		info:  log.New(os.Stderr, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile),
 		warn:  log.New(os.Stderr, "WARNING: ", log.Ldate|log.Ltime|log.Lshortfile),
 		error: log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile),
 	}
+
+	var nets []IpNetType
+	nets = append(nets, netTypes...)
+	if len(nets) == 0 {
+		nets = append(nets, IP64)
+	}
+
+	var dialer net.Dialer
+	np := newNetProbe(nets)
+	transport := http.DefaultTransport.(*http.Transport).Clone()
+	transport.DialTLSContext = nil
+	transport.DialContext, _ = np.NextNetworkType(dialer)
+
 	return Bilibili{
 		loggerCommon: logger,
 		userAgent:    kUserAgent,
 		http:         http.DefaultClient,
 		ctx:          ctx,
+		netTypes:     nets,
 	}
 }
 
-func NewBilibili() Bilibili {
+func NewBilibiliWithNetType(netTypes []IpNetType) Bilibili {
 	ctx := context.Background()
-	return NewBilibiliWithContext(ctx)
+	return NewBilibiliWithContext(ctx, netTypes)
+}
+
+func NewBilibili() Bilibili {
+	return NewBilibiliWithNetType(nil)
 }
diff --git a/bilibili/netprobe.go b/bilibili/netprobe.go
new file mode 100644
index 0000000..ac11a77
--- /dev/null
+++ b/bilibili/netprobe.go
@@ -0,0 +1,41 @@
+package bilibili
+
+import (
+	"context"
+	"net"
+)
+
+type IpNetType string
+
+var (
+	IPv6Net IpNetType = "tcp6"
+	IPv4Net IpNetType = "tcp4"
+	IP64    IpNetType = "tcp"
+)
+
+type netContext = func(context.Context, string, string) (net.Conn, error)
+
+type netProbe struct {
+	list []IpNetType
+	i    int
+}
+
+func newNetProbe(protocols []IpNetType) netProbe {
+	var netList []IpNetType
+	netList = append(netList, protocols...)
+	return netProbe{
+		list: netList,
+		i:    0,
+	}
+}
+
+func (p *netProbe) NextNetworkType(dialer net.Dialer) (netContext, IpNetType) {
+	if p.i >= len(p.list) {
+		return nil, IP64
+	}
+	network := p.list[p.i]
+	p.i++
+	return func(ctx context.Context, _, addr string) (net.Conn, error) {
+		return dialer.DialContext(ctx, string(network), addr)
+	}, network
+}
diff --git a/bilibili/request.go b/bilibili/request.go
index c31cf75..f9e2088 100644
--- a/bilibili/request.go
+++ b/bilibili/request.go
@@ -1,9 +1,11 @@
 package bilibili
 
 import (
+	"bilibili-livestream-archiver/common"
 	"encoding/json"
 	"io"
 	"log"
+	"net"
 	"net/http"
 	"strings"
 )
@@ -37,7 +39,7 @@ func callGet[T BaseResponse[V], V any](b Bilibili, url string) (resp T, err erro
 		return
 	}
 
-	r, err := b.http.Do(req)
+	r, err := b.Do(req)
 	if err != nil {
 		logger.Printf("ERROR: HTTP Request failed on API %v: %v", url, err)
 		return
@@ -66,3 +68,25 @@ func callGet[T BaseResponse[V], V any](b Bilibili, url string) (resp T, err erro
 	b.debug.Printf("HTTP %v, len: %v bytes, url: %v", r.StatusCode, len(data), url)
 	return
 }
+
+func (b Bilibili) Do(req *http.Request) (resp *http.Response, err error) {
+	transport := http.DefaultTransport.(*http.Transport).Clone()
+	transport.DialTLSContext = nil
+
+	np := newNetProbe(b.netTypes)
+	var dialer net.Dialer
+	for netCtx, typeName := np.NextNetworkType(dialer); netCtx != nil; netCtx, typeName = np.NextNetworkType(dialer) {
+		transport.DialContext = netCtx
+		b.http.Transport = transport
+		resp, err = b.http.Do(req)
+
+		isOpErr := common.IsErrorOfType(err, &net.OpError{})
+		isAddrErr := common.IsErrorOfType(err, &net.AddrError{})
+		if err == nil || !isOpErr || !isAddrErr {
+			// return the first success request
+			b.loggerCommon.info.Printf("Request success with network %v.", typeName)
+			return
+		}
+	}
+	return
+}
diff --git a/bilibili/streaming.go b/bilibili/streaming.go
index 4ca8682..e8a6911 100644
--- a/bilibili/streaming.go
+++ b/bilibili/streaming.go
@@ -32,7 +32,7 @@ func (b Bilibili) CopyLiveStream(
 	r.Header.Set("Referer",
 		fmt.Sprintf("https://live.bilibili.com/blanc/%d?liteVersion=true", roomId))
 
-	resp, err := b.http.Do(r)
+	resp, err := b.Do(r)
 	if err != nil {
 		b.error.Printf("Cannot make HTTP GET request on %v: %v\n", url, err)
 		return
-- 
cgit v1.2.3