feat(driver): add CallActionWithContext (#107)

Signed-off-by: Me0wo <[email protected]>

Author: Sn0wo2

api.go

@@ -2,6 +2,7 @@ package zero
 
 import (
 	"bytes"
+	"context"
 	"crypto/md5"
 	"encoding/base64"
 	"encoding/hex"
@@ -10,6 +11,7 @@ import (
 	"fmt"
 	"regexp"
 	"strconv"
+	"time"
 
 	log "github.com/sirupsen/logrus"
 	"github.com/tidwall/gjson"
@@ -51,11 +53,17 @@ func formatMessage(msg interface{}) string {
 
 // CallAction 调用 cqhttp API
 func (ctx *Ctx) CallAction(action string, params Params) APIResponse {
-	req := APIRequest{
+	c, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+	return ctx.CallActionWithContext(c, action, params)
+}
+
+// CallActionWithContext 使用 context 调用 cqhttp API
+func (ctx *Ctx) CallActionWithContext(c context.Context, action string, params Params) APIResponse {
+	rsp, err := ctx.caller.CallAPI(c, APIRequest{
 		Action: action,
 		Params: params,
-	}
-	rsp, err := ctx.caller.CallAPI(req)
+	})
 	if err != nil {
 		log.Errorln("[api] 调用", action, "时出现错误: ", err)
 	}

bot.go

@@ -1,6 +1,7 @@
 package zero
 
 import (
+	"context"
 	"encoding/json"
 	"hash/crc64"
 	"runtime/debug"
@@ -37,7 +38,7 @@ var APICallers callerMap
 
 // APICaller is the interface of CallAPI
 type APICaller interface {
-	CallAPI(request APIRequest) (APIResponse, error)
+	CallAPI(c context.Context, request APIRequest) (APIResponse, error)
 }
 
 // Driver 与OneBot通信的驱动，使用driver.DefaultWebSocketDriver
@@ -137,14 +138,14 @@ type messageLogger struct {
 }
 
 // CallAPI 记录被触发的回复消息
-func (m *messageLogger) CallAPI(request APIRequest) (rsp APIResponse, err error) {
+func (m *messageLogger) CallAPI(ctx context.Context, request APIRequest) (rsp APIResponse, err error) {
 	noLog := false
 	b, ok := request.Params["__zerobot_no_log_mseeage_id__"].(bool)
 	if ok {
 		noLog = b
 		delete(request.Params, "__zerobot_no_log_mseeage_id__")
 	}
-	rsp, err = m.caller.CallAPI(request)
+	rsp, err = m.caller.CallAPI(ctx, request)
 	if err != nil {
 		return
 	}

driver/http.go

@@ -2,6 +2,7 @@ package driver
 
 import (
 	"bytes"
+	"context"
 	"crypto/hmac"
 	"crypto/sha1"
 	"encoding/hex"
@@ -30,7 +31,9 @@ type HTTP struct {
 
 func (h *HTTP) Connect() {
 	log.Infof("[httpcaller] 正在尝试与服务器握手: %s", h.caller.URL)
-	rsp, err := h.caller.CallAPI(zero.APIRequest{Action: "get_login_info", Params: nil})
+	c, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+	rsp, err := h.caller.CallAPI(c, zero.APIRequest{Action: "get_login_info", Params: nil})
 	if err != nil {
 		log.Warningf("[httpcaller] 与服务器握手失败: %s\n%v", h.caller.URL, err)
 		return
@@ -148,8 +151,8 @@ func (h *HTTP) Listen(handler func([]byte, zero.APICaller)) {
 
 // httpCaller 对 api 进行调用
 // 不关闭body会导致资源泄漏!
-func (c *HTTPCaller) httpCaller(action string, payload []byte) (*http.Response, error) {
-	req, err := http.NewRequest(http.MethodPost, c.URL+"/"+action, bytes.NewBuffer(payload))
+func (c *HTTPCaller) httpCaller(ctx context.Context, action string, payload []byte) (*http.Response, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.URL+"/"+action, bytes.NewBuffer(payload))
 	if err != nil {
 		return nil, err
 	}
@@ -169,13 +172,13 @@ func (c *HTTPCaller) httpCaller(action string, payload []byte) (*http.Response,
 	return resp, nil
 }
 
-func (c *HTTPCaller) CallAPI(request zero.APIRequest) (zero.APIResponse, error) {
-	p, err := json.Marshal(request.Params)
+func (c *HTTPCaller) CallAPI(ctx context.Context, req zero.APIRequest) (zero.APIResponse, error) {
+	p, err := json.Marshal(req.Params)
 	if err != nil {
 		return nullResponse, err
 	}
 
-	resp, err := c.httpCaller(request.Action, p)
+	resp, err := c.httpCaller(ctx, req.Action, p)
 	if err != nil {
 		return nullResponse, err
 	}

driver/wsclient.go

@@ -1,11 +1,11 @@
 package driver
 
 import (
+	"context"
 	"encoding/base64"
 	"io"
 	"net"
 	"net/http"
-	"os"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -142,7 +142,7 @@ func (ws *WSClient) nextSeq() uint64 {
 }
 
 // CallAPI 发送ws请求
-func (ws *WSClient) CallAPI(req zero.APIRequest) (zero.APIResponse, error) {
+func (ws *WSClient) CallAPI(c context.Context, req zero.APIRequest) (zero.APIResponse, error) {
 	ch := make(chan zero.APIResponse, 1)
 	req.Echo = ws.nextSeq()
 	ws.seqMap.Store(req.Echo, ch)
@@ -163,7 +163,7 @@ func (ws *WSClient) CallAPI(req zero.APIRequest) (zero.APIResponse, error) {
 			return nullResponse, io.ErrClosedPipe
 		}
 		return rsp, nil
-	case <-time.After(time.Minute):
-		return nullResponse, os.ErrDeadlineExceeded
+	case <-c.Done():
+		return nullResponse, c.Err()
 	}
 }

driver/wsserver.go

@@ -1,12 +1,12 @@
 package driver
 
 import (
+	"context"
 	"encoding/json"
 	"io"
 	"net"
 	"net/http"
 	"net/url"
-	"os"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -213,7 +213,7 @@ func (wssc *WSSCaller) nextSeq() uint64 {
 }
 
 // CallAPI 发送ws请求
-func (wssc *WSSCaller) CallAPI(req zero.APIRequest) (zero.APIResponse, error) {
+func (wssc *WSSCaller) CallAPI(c context.Context, req zero.APIRequest) (zero.APIResponse, error) {
 	ch := make(chan zero.APIResponse, 1)
 	req.Echo = wssc.nextSeq()
 	wssc.seqMap.Store(req.Echo, ch)
@@ -234,7 +234,7 @@ func (wssc *WSSCaller) CallAPI(req zero.APIRequest) (zero.APIResponse, error) {
 			return nullResponse, io.ErrClosedPipe
 		}
 		return rsp, nil
-	case <-time.After(time.Minute):
-		return nullResponse, os.ErrDeadlineExceeded
+	case <-c.Done():
+		return nullResponse, c.Err()
 	}
 }

pattern_test.go

@@ -1,17 +1,19 @@
 package zero
 
 import (
+	"context"
 	"fmt"
+	"strconv"
+	"testing"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/tidwall/gjson"
 	"github.com/wdvxdr1123/ZeroBot/message"
-	"strconv"
-	"testing"
 )
 
 type mockAPICaller struct{}
 
-func (m mockAPICaller) CallAPI(_ APIRequest) (APIResponse, error) {
+func (m mockAPICaller) CallAPI(_ context.Context, _ APIRequest) (APIResponse, error) {
 	return APIResponse{
 		Status:  "",
 		Data:    gjson.Parse(`{"message_id":"12345","sender":{"user_id":12345}}`), // just for reply cleaner