From 5fb5a55f79f0fb86d70cade132c2d256d9f806fc Mon Sep 17 00:00:00 2001 From: Jake Hillion Date: Sun, 25 Oct 2020 15:36:34 +0000 Subject: [PATCH] reordering, formatting, some tests --- .gitignore | 1 + config/builder.go | 16 +++++++------- config/loader.go | 4 ++-- go.mod | 1 + main.go | 12 ++++++----- proxy/packet.go | 4 ++-- proxy/packet_test.go | 50 ++++++++++++++++++++++++++++++++++++++++++++ proxy/proxy.go | 20 ++++++++++++------ tcp/flow.go | 8 +++---- tcp/receiver.go | 8 +++---- utils/utils.go | 2 +- 11 files changed, 94 insertions(+), 32 deletions(-) create mode 100644 proxy/packet_test.go diff --git a/.gitignore b/.gitignore index 0b09f86..5f53fd7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ config.ini Makefile +logs/ # Created by https://www.toptal.com/developers/gitignore/api/intellij+all,go # Edit at https://www.toptal.com/developers/gitignore?templates=intellij+all,go diff --git a/config/builder.go b/config/builder.go index 170e7aa..9cf31d1 100644 --- a/config/builder.go +++ b/config/builder.go @@ -26,6 +26,14 @@ func (c Configuration) Build() (*proxy.Proxy, error) { p := proxy.Proxy{} p.Generator = UselessMac(false) + ss, err := tun.NewTun("nc%d", 1500) + if err != nil { + return nil, err + } + + p.Source = ss + p.Sink = ss + for _, peer := range c.Peers { p.Verifiers = append(p.Verifiers, UselessMac(true)) @@ -40,14 +48,6 @@ func (c Configuration) Build() (*proxy.Proxy, error) { } } - ss, err := tun.NewTun("nc%d", 1500) - if err != nil { - return nil, err - } - - p.Source = ss - p.Sink = ss - return &p, nil } diff --git a/config/loader.go b/config/loader.go index 4bd55af..c4fdd3f 100644 --- a/config/loader.go +++ b/config/loader.go @@ -7,8 +7,8 @@ import ( func LoadConfig(path string) (c Configuration, err error) { var file *ini.File file, err = ini.LoadSources(ini.LoadOptions{ - AllowShadows: true, - AllowNonUniqueSections: true, + AllowShadows: true, + AllowNonUniqueSections: true, }, path) if err != nil { return diff --git a/go.mod b/go.mod index b2bad3d..2a9eb20 100644 --- a/go.mod +++ b/go.mod @@ -5,5 +5,6 @@ go 1.15 require ( github.com/go-playground/validator/v10 v10.4.1 github.com/pkg/taptun v0.0.0-20160424131934-bbbd335672ab + github.com/stretchr/testify v1.4.0 gopkg.in/ini.v1 v1.62.0 ) diff --git a/main.go b/main.go index adbe0de..d7f2633 100644 --- a/main.go +++ b/main.go @@ -1,7 +1,7 @@ package main import ( - "fmt" + "log" "mpbl3p/config" "os" "os/signal" @@ -9,23 +9,25 @@ import ( ) func main() { - fmt.Println("loading config...") + log.SetFlags(log.Ldate | log.Ltime | log.Llongfile) + + log.Println("loading config...") c, err := config.LoadConfig("config.ini") if err != nil { panic(err) } - fmt.Println("building config...") + log.Println("building config...") p, err := c.Build() if err != nil { panic(err) } - fmt.Println("starting...") + log.Println("starting...") p.Start() - fmt.Println("running") + log.Println("running") signals := make(chan os.Signal) signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT) diff --git a/proxy/packet.go b/proxy/packet.go index 66f6134..1588278 100644 --- a/proxy/packet.go +++ b/proxy/packet.go @@ -13,7 +13,7 @@ type Packet struct { // create a packet from the raw data of an IP packet func NewPacket(data []byte) Packet { return Packet{ - Data: data, + Data: data, timestamp: time.Now(), } } @@ -29,7 +29,7 @@ func UnmarshalPacket(raw []byte, verifier MacVerifier) (Packet, error) { } p := Packet{ - Data: data[:8], + Data: data[:len(data)-8], } unixTime := int64(binary.LittleEndian.Uint64(data[len(data)-8:])) diff --git a/proxy/packet_test.go b/proxy/packet_test.go new file mode 100644 index 0000000..ca48a7e --- /dev/null +++ b/proxy/packet_test.go @@ -0,0 +1,50 @@ +package proxy + +import ( + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "testing" +) + +type AlmostUselessMac bool + +func (AlmostUselessMac) CodeLength() int { + return 4 +} + +func (AlmostUselessMac) Generate([]byte) []byte { + return []byte{'a', 'b', 'c', 'd'} +} + +func (u AlmostUselessMac) Verify(_, sum []byte) error { + if !(sum[0] == 'a' && sum[1] == 'b' && sum[2] == 'c' && sum[3] == 'd') { + return ErrBadChecksum + } + return nil +} + +func TestPacket_Marshal(t *testing.T) { + testContent := []byte("A test string is the content of this packet.") + testPacket := NewPacket(testContent) + testMac := AlmostUselessMac(false) + + t.Run("Length", func(t *testing.T) { + marshalled := testPacket.Marshal(testMac) + + assert.Len(t, marshalled, len(testContent)+8+4) + }) +} + +func TestUnmarshalPacket(t *testing.T) { + testContent := []byte("A test string is the content of this packet.") + testPacket := NewPacket(testContent) + testMac := AlmostUselessMac(false) + testMarshalled := testPacket.Marshal(testMac) + + t.Run("Success", func(t *testing.T) { + content, err := UnmarshalPacket(testMarshalled, testMac) + + require.Nil(t, err) + assert.Equal(t, testContent, content.Raw()) + }) +} diff --git a/proxy/proxy.go b/proxy/proxy.go index 3e10143..5c10625 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -1,6 +1,8 @@ package proxy -import "fmt" +import ( + "log" +) type Producer interface { IsAlive() bool @@ -44,6 +46,7 @@ func (p *Proxy) Start() { panic(err) return } else { + log.Println("proxy received packet from source") p.proxyChan <- packet } } @@ -55,23 +58,26 @@ func (p *Proxy) Start() { go func() { for c.IsAlive() { if err := c.Consume(<-p.proxyChan, p.Generator); err != nil { - fmt.Println(err) + log.Println(err) return + } else { + log.Println("proxy consumed a packet") } } }() } for i, prod := range p.Producers { - i := i + verifier := p.Verifiers[i] prod := prod go func() { for prod.IsAlive() { - if packet, err := prod.Produce(p.Verifiers[i]); err != nil { - fmt.Println(err) + if packet, err := prod.Produce(verifier); err != nil { + log.Println(err) return } else { + log.Println("proxy received a packet from a producer") p.sinkChan <- packet } } @@ -84,7 +90,9 @@ func (p *Proxy) Start() { if err := p.Sink.Sink(packet); err != nil { panic(err) return + } else { + log.Println("proxy sent packet via sink") } } }() -} \ No newline at end of file +} diff --git a/tcp/flow.go b/tcp/flow.go index 9a06099..b373252 100644 --- a/tcp/flow.go +++ b/tcp/flow.go @@ -19,7 +19,7 @@ type Flow struct { func NewFlow(local, remote string) (Flow, error) { f := Flow{ - Local: local, + Local: local, Remote: remote, } @@ -67,7 +67,7 @@ func (f *Flow) Consume(p proxy.Packet, g proxy.MacGenerator) error { func (f *Flow) consumeMarshalled(data []byte) error { prefixedData := make([]byte, len(data)+4) binary.LittleEndian.PutUint32(prefixedData, uint32(len(data))) - copy(prefixedData[4:], data) + copy(prefixedData[:len(prefixedData)-4], data) _, err := f.conn.Write(prefixedData) return err @@ -82,7 +82,7 @@ func (f *Flow) Produce(v proxy.MacVerifier) (proxy.Packet, error) { return proxy.UnmarshalPacket(data, v) } -func (f *Flow) produceMarshalled() ([]byte,error) { +func (f *Flow) produceMarshalled() ([]byte, error) { lengthBytes := make([]byte, 4) if n, err := io.LimitReader(f.conn, 4).Read(lengthBytes); err != nil { return nil, err @@ -99,4 +99,4 @@ func (f *Flow) produceMarshalled() ([]byte,error) { } return dataBytes, nil -} \ No newline at end of file +} diff --git a/tcp/receiver.go b/tcp/receiver.go index bbe167d..4640380 100644 --- a/tcp/receiver.go +++ b/tcp/receiver.go @@ -1,7 +1,7 @@ package tcp import ( - "fmt" + "log" "mpbl3p/proxy" "mpbl3p/utils" "net" @@ -15,7 +15,7 @@ type Receiver struct { } func NewReceiver(local string) (*Receiver, error) { - r := Receiver{} + r := Receiver{flows: make(map[int]Flow)} laddr, err := net.ResolveTCPAddr("tcp", local) if err != nil { @@ -70,7 +70,7 @@ func (r *Receiver) addFlow(flow Flow) { for flow.IsAlive() { d, err := flow.produceMarshalled() if err != nil { - fmt.Println(err) + log.Println(err) continue } r.inPackets <- d @@ -83,7 +83,7 @@ func (r *Receiver) addFlow(flow Flow) { d := <-r.outPackets err := flow.consumeMarshalled(d) if err != nil { - fmt.Println(err) + log.Println(err) // continue } } diff --git a/utils/utils.go b/utils/utils.go index db0d401..2548b3b 100644 --- a/utils/utils.go +++ b/utils/utils.go @@ -10,4 +10,4 @@ func init() { i += 1 } }() -} \ No newline at end of file +}