feat(img.pool): add rkey cache

img/pool/hook.go

@@ -0,0 +1,31 @@
+package pool
+
+import (
+	"time"
+
+	"github.com/sirupsen/logrus"
+	zero "github.com/wdvxdr1123/ZeroBot"
+)
+
+func init() {
+	zero.OnMessage(zero.HasPicture).SetBlock(false).FirstPriority().Handle(func(ctx *zero.Ctx) {
+		img, ok := ctx.State["image_url"].([]string)
+		if !ok || len(img) == 0 {
+			return
+		}
+		if !ntcachere.MatchString(img[0]) { // is not NTQQ
+			return
+		}
+		rk, err := nturl(img[0]).rkey()
+		if err != nil {
+			logrus.Debugln("[imgpool] parse rkey error:", err, "image url:", img)
+			return
+		}
+		err = rs.set(time.Minute, rk)
+		if err != nil {
+			logrus.Debugln("[imgpool] set rkey error:", err, "rkey:", rk)
+			return
+		}
+		logrus.Debugln("[imgpool] set latest rkey:", rk)
+	})
+}

img/pool/img.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"regexp"
 	"strings"
+	"time"
 
 	"github.com/sirupsen/logrus"
 	"github.com/wdvxdr1123/ZeroBot/message"
@@ -88,7 +89,11 @@ func (m *Image) String() string {
 	if oldimgre.MatchString(m.item.u) {
 		return fmt.Sprintf(cacheurl, m.item.u)
 	}
-	nu, err := unpack(m.item.u)
+	rk, err := rs.rkey(time.Minute)
+	if err != nil {
+		logrus.Debugln("[imgpool] get reky error:", err)
+	}
+	nu, err := unpack(m.item.u, rk)
 	if err != nil {
 		return m.f
 	}
@@ -121,7 +126,8 @@ func (m *Image) Push(send ctxext.NoCtxSendMsg, get ctxext.NoCtxGetMsg) (hassent
 			u := e.Data["url"]
 			if ntcachere.MatchString(u) { // is NTQQ
 				raw := ""
-				raw, err = nturl(u).pack()
+				ntu := nturl(u)
+				raw, err = ntu.pack()
 				if err != nil {
 					logrus.Errorln("[imgpool] pack nturl err:", err)
 					err = nil
@@ -138,6 +144,19 @@ func (m *Image) Push(send ctxext.NoCtxSendMsg, get ctxext.NoCtxGetMsg) (hassent
 				if err != nil {
 					logrus.Errorln("[imgpool] item.push err:", err)
 					err = nil
+					return
+				}
+				raw, err = ntu.rkey()
+				if err != nil {
+					logrus.Errorln("[imgpool] parse rkey err:", err)
+					err = nil
+					return
+				}
+				err = rs.set(time.Minute, raw)
+				if err != nil {
+					logrus.Errorln("[imgpool] set rkey err:", err)
+					err = nil
+					return
 				}
 				return
 			}

img/pool/nt.go

@@ -26,7 +26,7 @@ var (
 
 type nturl string
 
-func unpack(raw string) (nturl, error) {
+func unpack(raw, rkey string) (nturl, error) {
 	if len(raw) != ntrawlen {
 		return "", ErrInvalidNTRaw
 	}
@@ -37,12 +37,14 @@ func unpack(raw string) (nturl, error) {
 		return "", ErrInvalidNTRaw
 	}
 	fileid = fileid[:b]
-	rkey := base64.RawURLEncoding.EncodeToString(rb[60:])
-	b = rb[ntrawlen-1]
-	if len(rkey) < int(b) {
-		return "", ErrInvalidNTRaw
+	if rkey == "" {
+		rkey = base64.RawURLEncoding.EncodeToString(rb[60:])
+		b = rb[ntrawlen-1]
+		if len(rkey) < int(b) {
+			return "", ErrInvalidNTRaw
+		}
+		rkey = rkey[:b]
 	}
-	rkey = rkey[:b]
 	return nturl(fmt.Sprintf(ntcacheurl, fileid, rkey)), nil
 }
 
@@ -67,3 +69,12 @@ func (nu nturl) pack() (string, error) {
 	buf[ntrawlen-1] = byte(len(rkey))
 	return binary.BytesToString(buf[:]), nil
 }
+
+// rkey get the embeded rkey
+func (nu nturl) rkey() (string, error) {
+	subs := ntcachere.FindStringSubmatch(string(nu))
+	if len(subs) != 3 {
+		return "", ErrInvalidNTURL
+	}
+	return subs[2], nil
+}

img/pool/nt_test.go

@@ -18,7 +18,7 @@ func TestNTPacker(t *testing.T) {
 		t.Fatal(err)
 	}
 	t.Log(hex.EncodeToString(binary.StringToBytes(raw)))
-	upknt, err := unpack(raw)
+	upknt, err := unpack(raw, "")
 	if err != nil {
 		t.Fatal(err)
 	}

img/pool/pool.go

@@ -40,6 +40,22 @@ func getItem(name string) (*item, error) {
 	return &item{name: name, u: u}, nil
 }
 
+// update 同步 item 为服务器最新
+func (t *item) update() error {
+	reg := registry.NewRegReader("reilia.fumiama.top:35354", "", "fumiama")
+	err := reg.ConnectIn(time.Second * 4)
+	if err != nil {
+		return err
+	}
+	u, err := reg.Get(t.name)
+	defer reg.Close()
+	if err != nil {
+		return err
+	}
+	t.u = u
+	return nil
+}
+
 // push 推送 item
 func (t *item) push(key string) (err error) {
 	for i := 0; i < 8; i++ {

img/pool/rkey.go

@@ -0,0 +1,41 @@
+package pool
+
+import "time"
+
+const rkeykey = "__latest_rkey__"
+
+var rs *rkeystorage
+
+func init() {
+	var err error
+	rs.item, err = newItem(rkeykey, "")
+	if err != nil {
+		panic(err)
+	}
+}
+
+type rkeystorage struct {
+	*item
+	lastrefresh time.Time
+}
+
+func (rs *rkeystorage) rkey(timeout time.Duration) (string, error) {
+	if time.Since(rs.lastrefresh) < timeout {
+		return rs.u, nil
+	}
+	err := rs.item.update()
+	if err != nil {
+		return "", err
+	}
+	rs.lastrefresh = time.Now()
+	return rs.u, nil
+}
+
+func (rs *rkeystorage) set(timeout time.Duration, rkey string) error {
+	if time.Since(rs.lastrefresh) < timeout { // 降低设置频次
+		return nil
+	}
+	rs.item.u = rkey
+	rs.lastrefresh = time.Now()
+	return rs.item.push("minamoto")
+}