Support specify IP network type (ipv4, ipv6, in arbitrary combination and priority)

7 files changed, 129 insertions, 10 deletions

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
diff --git a/common/errors.go b/common/errors.go
new file mode 100644
index 0000000..ba686f1
--- /dev/null
+++ b/common/errors.go
@@ -0,0 +1,29 @@
+package common
+
+import (
+	"errors"
+	"reflect"
+)
+
+// IsErrorOfType is a modified version of errors.Is, which loosen the check condition
+func IsErrorOfType(err, target error) bool {
+	if target == nil {
+		return err == target
+	}
+
+	isComparable := reflect.TypeOf(target).Comparable()
+	for {
+		if isComparable && reflect.TypeOf(target) == reflect.TypeOf(err) {
+			return true
+		}
+		if x, ok := err.(interface{ Is(error) bool }); ok && x.Is(target) {
+			return true
+		}
+		// TODO: consider supporting target.Is(err). This would allow
+		// user-definable predicates, but also may allow for coping with sloppy
+		// APIs, thereby making it easier to get away with them.
+		if err = errors.Unwrap(err); err == nil {
+			return false
+		}
+	}
+}
diff --git a/recording/config.go b/recording/config.go
index 6d8704b..c38d73f 100644
--- a/recording/config.go
+++ b/recording/config.go
@@ -1,6 +1,7 @@
 package recording
 
 import (
+	"bilibili-livestream-archiver/bilibili"
 	"bilibili-livestream-archiver/common"
 	"fmt"
 )
@@ -13,9 +14,10 @@ type TaskConfig struct {
 }
 
 type TransportConfig struct {
-	SocketTimeoutSeconds int `mapstructure:"socket_timeout_seconds"`
-	RetryIntervalSeconds int `mapstructure:"retry_interval_seconds"`
-	MaxRetryTimes        int `mapstructure:"max_retry_times"`
+	SocketTimeoutSeconds int                  `mapstructure:"socket_timeout_seconds"`
+	RetryIntervalSeconds int                  `mapstructure:"retry_interval_seconds"`
+	MaxRetryTimes        int                  `mapstructure:"max_retry_times"`
+	AllowedNetworkTypes  []bilibili.IpNetType `mapstructure:"allowed_network_types"`
 }
 
 type DownloadConfig struct {
diff --git a/recording/runner.go b/recording/runner.go
index 4fd9b46..974fa0e 100644
--- a/recording/runner.go
+++ b/recording/runner.go
@@ -42,7 +42,9 @@ func RunTask(ctx context.Context, wg *sync.WaitGroup, task *TaskConfig) {
 // doTask do the actual work, but returns synchronously.
 func doTask(ctx context.Context, task *TaskConfig) error {
 	logger := log.Default()
-	bi := bilibili.NewBilibili()
+	netTypes := task.Transport.AllowedNetworkTypes
+	logger.Printf("Network types: %v", netTypes)
+	bi := bilibili.NewBilibiliWithNetType(netTypes)
 	logger.Printf("Start task: room %v", task.RoomId)
 
 	authKey, url, err := getStreamingServer(task, logger, bi)
@@ -143,6 +145,7 @@ func record(
 		return
 	}
 
+	logger.Printf("INFO: Getting stream url...")
 	urlInfo, err := common.AutoRetry(
 		ctx,
 		func() (bilibili.RoomUrlInfoResponse, error) {