fumiama/go-registry
1
0
Fork 0
mirror of https://github.com/fumiama/go-registry.git synced 2026-06-04 23:40:27 +08:00
Code Activity

优化

Browse Source
This commit is contained in:
源文雨
2022-10-28 11:25:34 +08:00
parent 680755130a
commit 42411b3d58
7 changed files with 96 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
cmd.go
View File

@@ -2,10 +2,12 @@ package registry
import (
"crypto/md5"
"encoding/binary"
"errors"
"io"
"unsafe"
base14 "github.com/fumiama/go-base16384"
tea "github.com/fumiama/gofastTEA"
)
@@ -30,7 +32,7 @@ const (
)
var (
ErrMd5Mismatch = errors.New("cmdpacket.decrypt: md5 mismatch")
ErrMd5Mismatch = errors.New("cmd: md5 mismatch")
)
type CmdPacket struct {
@@ -112,6 +114,7 @@ func (c *CmdPacket) ReadFrom(f io.Reader) (n int64, err error) {
return int64(cnt), err
}
cnt, err = io.ReadFull(f, c.raw[:c.len])
cnt += 1 + 1 + 16
if err != nil {
return int64(cnt), err
}
@@ -133,7 +136,7 @@ func (c *CmdPacket) Write(buf []byte) (n int, err error) {
c.len = r.len
c.md5 = r.md5
copy(c.raw[:], r.raw[:c.len])
c.Data = nil
c.Data = c.Data[1+1+16+int(c.len):]
return 1 + 1 + 16 + int(c.len) - oldlen, nil
}
@@ -141,9 +144,9 @@ func (c *CmdPacket) Write(buf []byte) (n int, err error) {
func (c *CmdPacket) Encrypt(seq uint8) (raw []byte) {
setseq(c.t, seq)
c.len = uint8(c.t.EncryptLittleEndianTo(c.Data, sumtable, c.raw[:]))
(*slice)(unsafe.Pointer(&raw)).Data = unsafe.Pointer(&c.rawCmdPacket)
(*slice)(unsafe.Pointer(&raw)).Len = 1 + 1 + 16 + int(c.len)
(*slice)(unsafe.Pointer(&raw)).Cap = 1 + 1 + 16 + 255
(*slice)(unsafe.Pointer(&raw)).data = unsafe.Pointer(&c.rawCmdPacket)
(*slice)(unsafe.Pointer(&raw)).len = 1 + 1 + 16 + int(c.len)
(*slice)(unsafe.Pointer(&raw)).cap = 1 + 1 + 16 + 255
return
}
@@ -170,6 +173,19 @@ func setseq(t *tea.TEA, seq uint8) {
*(*uint8)(unsafe.Add(unsafe.Pointer(t), 15)) = seq
}
// randuint32 returns a lock free uint32 value.
//
//go:linkname randuint32 runtime.fastrand
func randuint32() uint32
//go:nosplit
func fill() []byte {
var b [8]byte
binary.LittleEndian.PutUint32(b[:4], randuint32())
binary.LittleEndian.PutUint32(b[4:8], randuint32())
return base14.Encode(b[:7])
}
// TEA encoding sumtable
var sumtable = [0x10]uint32{
0x9e3579b9,

40
cmd_test.go
View File

@@ -1,40 +0,0 @@
package registry
import (
"errors"
"testing"
)
func TestRegedit(t *testing.T) {
reg := NewRegedit("127.0.0.1:8888", "testpwd", "testsps", 127, 127)
err := reg.Connect()
if err != nil {
t.Fatal(err)
}
ret, err := reg.Get("test")
if err != nil && !errors.Is(err, ErrNoSuchKey) {
t.Fatal(err)
}
t.Log(ret)
if ret != "" {
err = reg.Del("test")
if err != nil {
t.Fatal(err)
}
}
err = reg.Set("test", "测试")
if err != nil {
t.Fatal(err)
}
ret, err = reg.Get("test")
if err != nil {
t.Fatal(err)
}
if ret != "测试" {
t.Fail()
}
err = reg.Close()
if err != nil {
t.Fatal(err)
}
}

8
go.mod
View File

@@ -2,6 +2,10 @@ module github.com/fumiama/go-registry
go 1.17
require github.com/fumiama/gofastTEA v0.0.10
require (
github.com/fumiama/go-base16384 v1.6.1
github.com/fumiama/go-simple-protobuf v0.1.0
github.com/fumiama/gofastTEA v0.0.10
)
require github.com/fumiama/go-simple-protobuf v0.1.0
require golang.org/x/text v0.3.7 // indirect

15
go.sum
View File

@@ -1,4 +1,19 @@
github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/fumiama/go-base16384 v1.6.1 h1:4yb4JgmBJDnQtq3XGXXdLrVwEnRpjhMUt4eAcsNeA30=
github.com/fumiama/go-base16384 v1.6.1/go.mod h1:OEn+947GV5gsbTAnyuUW/SrfxJYUdYupSIQXOuGOcXM=
github.com/fumiama/go-simple-protobuf v0.1.0 h1:rLzJgNqB6LHNDVMl81yyNt6ZKziWtVfu+ioF0edlEVw=
github.com/fumiama/go-simple-protobuf v0.1.0/go.mod h1:5yYNapXq1tQMOZg9bOIVhQlZk9pQqpuFIO4DZLbsdy4=
github.com/fumiama/gofastTEA v0.0.10 h1:JJJ+brWD4kie+mmK2TkspDXKzqq0IjXm89aGYfoGhhQ=
github.com/fumiama/gofastTEA v0.0.10/go.mod h1:RIdbYZyB4MbH6ZBlPymRaXn3cD6SedlCu5W/HHfMPBk=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

12
helper.go
View File

@@ -9,9 +9,9 @@ import "unsafe"
// Unlike reflect.SliceHeader, its Data field is sufficient to guarantee the
// data it references will not be garbage collected.
type slice struct {
Data unsafe.Pointer
Len int
Cap int
data unsafe.Pointer
len int
cap int
}
// BytesToString 没有内存开销的转换
@@ -23,8 +23,8 @@ func BytesToString(b []byte) string {
func StringToBytes(s string) (b []byte) {
bh := (*slice)(unsafe.Pointer(&b))
sh := (*slice)(unsafe.Pointer(&s))
bh.Data = sh.Data
bh.Len = sh.Len
bh.Cap = sh.Len
bh.data = sh.data
bh.len = sh.len
bh.cap = sh.len
return b
}

93
reg.go
View File

@@ -15,29 +15,28 @@ import (
)
var (
ErrGetKeyTooLong = errors.New("get key too long")
ErrDecAck = errors.New("decrypt ack error")
ErrInternalServer = errors.New("internal server error")
ErrPermissionDenied = errors.New("permission denied")
ErrSetKeyTooLong = errors.New("set key too long")
ErrSetValTooLong = errors.New("set val too long")
ErrUnknownAck = errors.New("unknown ack error")
ErrNoSuchKey = errors.New("no such key")
ErrRawDataTooLong = errors.New("raw data too long")
ErrMd5NotEqual = errors.New("md5 not equal")
ErrGetKeyTooLong = errors.New("reg: get key too long")
ErrDecAck = errors.New("reg: decrypt ack error")
ErrInternalServer = errors.New("reg: internal server error")
ErrPermissionDenied = errors.New("reg: permission denied")
ErrSetKeyTooLong = errors.New("reg: set key too long")
ErrSetValTooLong = errors.New("reg: set val too long")
ErrUnknownAck = errors.New("reg: unknown ack error")
ErrNoSuchKey = errors.New("reg: no such key")
ErrRawDataTooLong = errors.New("reg: raw data too long")
ErrMd5NotEqual = errors.New("reg: md5 not equal")
)
type Regedit struct {
sync.Mutex
conn net.Conn
addr string
tp tea.TEA
ts *tea.TEA
dksz, ddsz int
seq byte
mu sync.Mutex
conn net.Conn
addr string
tp tea.TEA
ts *tea.TEA
seq byte
}
func NewRegedit(addr, pwd, sps string, dksz, ddsz int) *Regedit {
func NewRegedit(addr, pwd, sps string) *Regedit {
var tp, ts [16]byte
if len(pwd) > 15 {
pwd = pwd[:15]
@@ -45,44 +44,44 @@ func NewRegedit(addr, pwd, sps string, dksz, ddsz int) *Regedit {
if len(sps) > 15 {
sps = sps[:15]
}
copy(tp[:], StringToBytes(pwd))
copy(ts[:], StringToBytes(sps))
copy(tp[:], pwd)
copy(ts[:], sps)
s := tea.NewTeaCipherLittleEndian(ts[:])
return &Regedit{addr: addr, tp: tea.NewTeaCipherLittleEndian(tp[:]), ts: &s, dksz: dksz, ddsz: ddsz}
return &Regedit{addr: addr, tp: tea.NewTeaCipherLittleEndian(tp[:]), ts: &s}
}
func NewRegReader(addr, pwd string, dksz, ddsz int) *Regedit {
func NewRegReader(addr, pwd string) *Regedit {
var tp [16]byte
if len(pwd) > 15 {
pwd = pwd[:15]
}
copy(tp[:], StringToBytes(pwd))
return &Regedit{addr: addr, tp: tea.NewTeaCipherLittleEndian(tp[:]), dksz: dksz, ddsz: ddsz}
copy(tp[:], pwd)
return &Regedit{addr: addr, tp: tea.NewTeaCipherLittleEndian(tp[:])}
}
func (r *Regedit) Connect() (err error) {
r.Lock()
r.mu.Lock()
if r.conn == nil {
r.conn, err = net.Dial("tcp", r.addr)
}
r.Unlock()
r.mu.Unlock()
return
}
func (r *Regedit) ConnectIn(wait time.Duration) (err error) {
r.Lock()
func (r *Regedit) ConnectIn(timeout time.Duration) (err error) {
r.mu.Lock()
if r.conn == nil {
r.conn, err = net.DialTimeout("tcp", r.addr, wait)
r.conn, err = net.DialTimeout("tcp", r.addr, timeout)
}
r.Unlock()
r.mu.Unlock()
return
}
func (r *Regedit) Close() (err error) {
r.Lock()
defer r.Unlock()
r.mu.Lock()
defer r.mu.Unlock()
if r.conn != nil {
p := NewCmdPacket(CMDEND, []byte("fill"), &r.tp)
p := NewCmdPacket(CMDEND, fill(), &r.tp)
r.conn.Write(p.Encrypt(r.seq))
p.Put()
r.seq = 0
@@ -99,17 +98,17 @@ func (r *Regedit) Get(key string) (string, error) {
}
p := NewCmdPacket(CMDGET, StringToBytes(key), &r.tp)
defer p.Put()
r.Lock()
r.mu.Lock()
r.conn.Write(p.Encrypt(r.seq))
r.seq++
err := r.ack(p)
if err != nil {
r.Unlock()
r.mu.Unlock()
return "", err
}
err = p.Decrypt(r.seq)
r.seq++
r.Unlock()
r.mu.Unlock()
if err != nil {
return "", ErrDecAck
}
@@ -124,14 +123,14 @@ func (r *Regedit) Get(key string) (string, error) {
}
func (r *Regedit) Cat() (*Storage, error) {
p := NewCmdPacket(CMDCAT, []byte("fill"), &r.tp)
p := NewCmdPacket(CMDCAT, fill(), &r.tp)
defer p.Put()
r.Lock()
r.mu.Lock()
r.conn.Write(p.Encrypt(r.seq))
r.seq++
seq := r.seq
r.seq++
r.Unlock()
r.mu.Unlock()
var buf [64]byte
i := 0
for {
@@ -179,17 +178,17 @@ func (r *Regedit) Cat() (*Storage, error) {
func (r *Regedit) IsMd5Equal(m [md5.Size]byte) (bool, error) {
p := NewCmdPacket(CMDMD5, m[:], &r.tp)
defer p.Put()
r.Lock()
r.mu.Lock()
r.conn.Write(p.Encrypt(r.seq))
r.seq++
err := r.ack(p)
if err != nil {
r.Unlock()
r.mu.Unlock()
return false, err
}
err = p.Decrypt(r.seq)
r.seq++
r.Unlock()
r.mu.Unlock()
if err != nil {
return false, ErrDecAck
}
@@ -218,8 +217,8 @@ func (r *Regedit) Set(key, value string) error {
}
p := NewCmdPacket(CMDSET, StringToBytes(key), r.ts)
defer p.Put()
r.Lock()
defer r.Unlock()
r.mu.Lock()
defer r.mu.Unlock()
r.conn.Write(p.Encrypt(r.seq))
r.seq++
ack := NewCmdPacket(CMDACK, nil, &r.tp)
@@ -271,19 +270,19 @@ func (r *Regedit) Del(key string) error {
}
p := NewCmdPacket(CMDDEL, StringToBytes(key), r.ts)
defer p.Put()
r.Lock()
r.mu.Lock()
r.conn.Write(p.Encrypt(r.seq))
r.seq++
ack := NewCmdPacket(CMDACK, nil, &r.tp)
defer ack.Put()
err := r.ack(ack)
if err != nil {
r.Unlock()
r.mu.Unlock()
return err
}
err = ack.Decrypt(r.seq)
r.seq++
r.Unlock()
r.mu.Unlock()
if err != nil {
return ErrDecAck
}

4
reg_test.go
View File

@@ -5,7 +5,7 @@ import (
)
func TestReg(t *testing.T) {
r := NewRegedit("127.0.0.1:8888", "testpwd", "testsps", 127, 127)
r := NewRegedit("127.0.0.1:8888", "testpwd", "testsps")
err := r.Connect()
if err != nil {
t.Fatal(err)
@@ -67,7 +67,7 @@ func TestReg(t *testing.T) {
}
func TestPush(t *testing.T) {
r := NewRegedit("reilia.fumiama.top:32664", "fumiama", "--", 127, 127)
r := NewRegedit("reilia.fumiama.top:32664", "fumiama", "--")
err := r.Connect()
if err != nil {
t.Fatal(err)