Skip to content

Commit

Permalink
Support cat msg v2 (#56)
Browse files Browse the repository at this point in the history
* Support cat msg v2

* fixes

* Fixes 2

* Fix unit

* handle settings in storage cat interactions

* Fxies
  • Loading branch information
reshke authored Aug 26, 2024
1 parent 183dd69 commit a6039d4
Show file tree
Hide file tree
Showing 15 changed files with 290 additions and 90 deletions.
3 changes: 2 additions & 1 deletion cmd/client/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import (
"github.com/yezzey-gp/yproxy/pkg/message"
"github.com/yezzey-gp/yproxy/pkg/object"
"github.com/yezzey-gp/yproxy/pkg/proc"
"github.com/yezzey-gp/yproxy/pkg/settings"
"github.com/yezzey-gp/yproxy/pkg/tablespace"
"github.com/yezzey-gp/yproxy/pkg/ylogger"
)
Expand Down Expand Up @@ -105,7 +106,7 @@ func putFunc(con net.Conn, instanceCnf *config.Instance, args []string) error {
ycl := client.NewYClient(con)
r := proc.NewProtoReader(ycl)

msg := message.NewPutMessageV2(args[0], encrypt, []message.PutSettings{
msg := message.NewPutMessageV2(args[0], encrypt, []settings.StorageSettings{
{
Name: message.StorageClassSetting,
Value: storageClass,
Expand Down
88 changes: 88 additions & 0 deletions pkg/message/cat_message_v2.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
package message

import (
"encoding/binary"

"github.com/yezzey-gp/yproxy/pkg/settings"
)

type CatMessageV2 struct {
Decrypt bool
Name string
StartOffset uint64

Settings []settings.StorageSettings
}

var _ ProtoMessage = &CatMessage{}

func NewCatMessageV2(name string, decrypt bool, StartOffset uint64, Settings []settings.StorageSettings) *CatMessageV2 {
return &CatMessageV2{
Name: name,
Decrypt: decrypt,
StartOffset: StartOffset,
Settings: Settings,
}
}

func (c *CatMessageV2) Encode() []byte {
bt := []byte{
byte(MessageTypeCatV2),
0,
0,
0,
}

if c.Decrypt {
bt[1] = byte(DecryptMessage)
} else {
bt[1] = byte(NoDecryptMessage)
}

bt = append(bt, []byte(c.Name)...)
bt = append(bt, 0)
bt = binary.BigEndian.AppendUint64(bt, c.StartOffset)

bt = binary.BigEndian.AppendUint64(bt, uint64(len(c.Settings)))

for _, s := range c.Settings {

bt = append(bt, []byte(s.Name)...)
bt = append(bt, 0)

bt = append(bt, []byte(s.Value)...)
bt = append(bt, 0)
}

ln := len(bt) + 8

bs := make([]byte, 8)
binary.BigEndian.PutUint64(bs, uint64(ln))
return append(bs, bt...)
}

func (c *CatMessageV2) Decode(body []byte) {
var off uint64
c.Name, off = GetCstring(body[4:])
if body[1] == byte(DecryptMessage) {
c.Decrypt = true
}
c.StartOffset = binary.BigEndian.Uint64(body[4+len(c.Name)+1:])

settLen := binary.BigEndian.Uint64(body[4+8+off : 4+off+8+8])

totalOff := 4 + off + 8 + 8

c.Settings = make([]settings.StorageSettings, settLen)

for i := 0; i < int(settLen); i++ {

var currOff uint64

c.Settings[i].Name, currOff = GetCstring(body[totalOff:])
totalOff += currOff

c.Settings[i].Value, currOff = GetCstring(body[totalOff:])
totalOff += currOff
}
}
3 changes: 3 additions & 0 deletions pkg/message/message.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ const (
MessageTypeCopy = MessageType(51)
MessageTypeGool = MessageType(52)
MessageTypePutV2 = MessageType(53)
MessageTypeCatV2 = MessageType(54)

DecryptMessage = RequestEncryption(1)
NoDecryptMessage = RequestEncryption(0)
Expand All @@ -36,6 +37,8 @@ func (m MessageType) String() string {
switch m {
case MessageTypeCat:
return "CAT"
case MessageTypeCatV2:
return "CATV2"
case MessageTypePut:
return "PUT"
case MessageTypePutV2:
Expand Down
67 changes: 65 additions & 2 deletions pkg/message/message_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (

"github.com/stretchr/testify/assert"
"github.com/yezzey-gp/yproxy/pkg/message"
"github.com/yezzey-gp/yproxy/pkg/settings"
)

func TestCatMsg(t *testing.T) {
Expand Down Expand Up @@ -81,15 +82,15 @@ func TestPutV2Msg(t *testing.T) {
name string
encrypt bool
err error
settings []message.PutSettings
settings []settings.StorageSettings
}

for _, tt := range []tcase{
{
"nam1",
true,
nil,
[]message.PutSettings{
[]settings.StorageSettings{
{
Name: "a",
Value: "b",
Expand All @@ -111,6 +112,68 @@ func TestPutV2Msg(t *testing.T) {

assert.Equal(msg.Name, msg2.Name)
assert.Equal(msg.Encrypt, msg2.Encrypt)
assert.Equal(msg.Settings, msg2.Settings)
}
}

func TestCatMsgV2(t *testing.T) {
assert := assert.New(t)

type tcase struct {
name string
decrypt bool
off uint64

settings []settings.StorageSettings
err error
}

for _, tt := range []tcase{
{
"nam1",
true,
0,
[]settings.StorageSettings{
{
Name: "a",
Value: "b",
},
{
Name: "cdsdsd",
Value: "ds",
},
},
nil,
},
{
"nam1",
true,
10,
[]settings.StorageSettings{
{
Name: "a",
Value: "b",
},
{
Name: "cdsdsd",
Value: "ds",
},
},
nil,
},
} {

msg := message.NewCatMessageV2(tt.name, tt.decrypt, tt.off, tt.settings)
body := msg.Encode()

msg2 := message.CatMessageV2{}

msg2.Decode(body[8:])

assert.Equal(msg.Name, msg2.Name)
assert.Equal(msg.Decrypt, msg2.Decrypt)
assert.Equal(msg.StartOffset, msg2.StartOffset)
assert.Equal(msg.Settings, msg2.Settings)
}
}

Expand Down
35 changes: 8 additions & 27 deletions pkg/message/put_message_v2.go
Original file line number Diff line number Diff line change
@@ -1,28 +1,24 @@
package message

import (
"bytes"
"encoding/binary"

"github.com/yezzey-gp/yproxy/pkg/settings"
)

const StorageClassSetting = "StorageClass"
const TableSpaceSetting = "TableSpace"

type PutSettings struct {
Name string
Value string
}

type PutMessageV2 struct {
Encrypt bool
Name string

Settings []PutSettings
Settings []settings.StorageSettings
}

var _ ProtoMessage = &PutMessageV2{}

func NewPutMessageV2(name string, encrypt bool, settings []PutSettings) *PutMessageV2 {
func NewPutMessageV2(name string, encrypt bool, settings []settings.StorageSettings) *PutMessageV2 {
return &PutMessageV2{
Name: name,
Encrypt: encrypt,
Expand Down Expand Up @@ -67,42 +63,27 @@ func (c *PutMessageV2) Encode() []byte {
return append(bs, bt...)
}

func (c *PutMessageV2) GetCstring(b []byte) (string, uint64) {
offset := uint64(0)
buff := bytes.NewBufferString("")

for i := 0; i < len(b); i++ {
offset++
if b[i] == 0 {
break
}
buff.WriteByte(b[i])
}

return buff.String(), offset
}

func (c *PutMessageV2) Decode(body []byte) {
if body[1] == byte(EncryptMessage) {
c.Encrypt = true
}
var off uint64
c.Name, off = c.GetCstring(body[4:])
c.Name, off = GetCstring(body[4:])

settLen := binary.BigEndian.Uint64(body[4+off : 4+off+8])

totalOff := 4 + off + 8

c.Settings = make([]PutSettings, settLen)
c.Settings = make([]settings.StorageSettings, settLen)

for i := 0; i < int(settLen); i++ {

var currOff uint64

c.Settings[i].Name, currOff = c.GetCstring(body[totalOff:])
c.Settings[i].Name, currOff = GetCstring(body[totalOff:])
totalOff += currOff

c.Settings[i].Value, currOff = c.GetCstring(body[totalOff:])
c.Settings[i].Value, currOff = GetCstring(body[totalOff:])
totalOff += currOff
}
}
18 changes: 18 additions & 0 deletions pkg/message/utils.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
package message

import "bytes"

func GetCstring(b []byte) (string, uint64) {
offset := uint64(0)
buff := bytes.NewBufferString("")

for i := 0; i < len(b); i++ {
offset++
if b[i] == 0 {
break
}
buff.WriteByte(b[i])
}

return buff.String(), offset
}
22 changes: 11 additions & 11 deletions pkg/mock/storage.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit a6039d4

Please sign in to comment.