diff --git a/core/iovec/iovec.go b/core/iovec/iovec.go index 7700aac..eacc4b8 100644 --- a/core/iovec/iovec.go +++ b/core/iovec/iovec.go @@ -33,6 +33,9 @@ func (self IoVec) Concat() []byte { } func (self *IoVec) Take(s []byte) *IoVec { + if len(s) == 0 { + return self + } *self = append(*self, s) return self } @@ -69,6 +72,31 @@ func (self *IoVec) Consume() []byte { return b.Bytes() } +func (self IoVec) LastByte() (byte, error) { + l := len(self) + if l == 0 { + return 0, fmt.Errorf("This IoVec is empty") + } + k := len(self[l-1]) + return self[l-1][k-1], nil +} + +func (self *IoVec) Drop(s int) error { + l := len(*self) + c := 0 + for i := l - 1; i >= 0; i-- { + v := (*self)[i] + vl := len(v) + if c+vl >= s { + (*self) = (*self)[:i] + self.Take(v[:c+vl-s]) + return nil + } + c += vl + } + return fmt.Errorf("Unable to drop %d bytes", s) +} + func (self IoVec) At(i int) (byte, error) { c := 0 for _, v := range self { @@ -91,17 +119,19 @@ func (self IoVec) Peek(i int) ([]byte, error) { return nil, fmt.Errorf("Index %d out of bound", i) } -func (self *IoVec) Drop(i int) error { +func (self *IoVec) Split(i int) (tail IoVec) { c := 0 for k, v := range *self { if i >= c && i < c+len(v) { - *self = (*self)[:k] - if i-c != 0 { - *self = append(*self, v[:i-c]) + tail.Take(v[i-c:]) + if len((*self)[k+1:]) != 0 { + tail = append(tail, (*self)[k+1:]...) } - return nil + *self = (*self)[:k] + self.Take(v[:i-c]) + return } c += len(v) } - return fmt.Errorf("Index %d out of bound", i) + return } diff --git a/core/iovec/iovec_test.go b/core/iovec/iovec_test.go index 223c62e..605d6ba 100644 --- a/core/iovec/iovec_test.go +++ b/core/iovec/iovec_test.go @@ -32,7 +32,7 @@ func TestDrop(t *testing.T) { var v IoVec v.Take([]byte("hello")) v.Take([]byte("bar")) - if err := v.Drop(5); err != nil { + if err := v.Drop(3); err != nil { t.Fail() } s := string(v.Consume()) @@ -45,7 +45,7 @@ func TestDrop2(t *testing.T) { var v IoVec v.Take([]byte("hello")) v.Take([]byte("bar")) - if err := v.Drop(6); err != nil { + if err := v.Drop(2); err != nil { t.Fail() } s := string(v.Consume()) @@ -54,6 +54,66 @@ func TestDrop2(t *testing.T) { } } +func TestSplit(t *testing.T) { + var v IoVec + v.Take([]byte("hello")) + v.Take([]byte("bar")) + tail := v.Split(0) + s := string(tail.Consume()) + if s != "hellobar" { + t.Fail() + } + sv := string(v.Consume()) + if sv != "" { + t.Fail() + } +} + +func TestSplit1(t *testing.T) { + var v IoVec + v.Take([]byte("hello")) + v.Take([]byte("bar")) + tail := v.Split(5) + s := string(tail.Consume()) + if s != "bar" { + t.Fail() + } + sv := string(v.Consume()) + if sv != "hello" { + t.Fail() + } +} + +func TestSplit2(t *testing.T) { + var v IoVec + v.Take([]byte("hello")) + v.Take([]byte("bar")) + tail := v.Split(6) + s := string(tail.Consume()) + if s != "ar" { + t.Fail() + } + sv := string(v.Consume()) + if sv != "hellob" { + t.Fail() + } +} + +func TestSplit3(t *testing.T) { + var v IoVec + v.Take([]byte("hello")) + v.Take([]byte("bar")) + tail := v.Split(3) + s := string(tail.Consume()) + if s != "lobar" { + t.Fail() + } + sv := string(v.Consume()) + if sv != "hel" { + t.Fail() + } +} + func TestConcat(t *testing.T) { var v IoVec v.Take([]byte("hello")) diff --git a/passes/passes.go b/passes/passes.go index aa5b170..ef1f3dc 100644 --- a/passes/passes.go +++ b/passes/passes.go @@ -208,12 +208,11 @@ type OBFSDecoder struct { } func (self *OBFSDecoder) Run(b *iovec.IoVec) error { - l := b.Len() - t, err := b.At(l - 1) + t, err := b.LastByte() if err != nil { return err } - if err := b.Drop(l - (1 + int(t))); err != nil { + if err := b.Drop(1 + int(t)); err != nil { return err } return WrapLegacyPass(self, b)