Merge pull request #9 from keyvchan/generic

refactor: message passing

Author: keyvchan

client/client.go

@@ -3,12 +3,14 @@ package client
 import (
 	"log"
 
-	"github.com/keyvchan/NetAssist/internal"
+	"github.com/keyvchan/NetAssist/pkg/flags"
+	"github.com/keyvchan/NetAssist/pkg/utils"
 	"github.com/keyvchan/NetAssist/protocol"
 )
 
+// Req is the entry point for the client
 func Req() {
-	types := internal.GetArg(2)
+	types := flags.GetArg(2)
 	log.Println("Req:", types)
 	switch types {
 	case "tcp":
@@ -20,9 +22,9 @@ func Req() {
 	case "unixgram":
 		protocol.UnixgramClient()
 	case "unixpacket":
-		internal.Unimplemented("unixpacket")
+		utils.Unimplemented("unixpacket")
 	case "ip":
-		internal.Unimplemented("ip")
+		utils.Unimplemented("ip")
 	default:
 		log.Fatal("unknown protocol", types)
 	}

internal/read.go

@@ -4,13 +4,13 @@ import (
 	"log"
 
 	"github.com/keyvchan/NetAssist/client"
-	"github.com/keyvchan/NetAssist/internal"
+	"github.com/keyvchan/NetAssist/pkg/flags"
 	"github.com/keyvchan/NetAssist/server"
 )
 
 func main() {
-	internal.SetArgs()
-	types := internal.GetArg(1)
+	flags.SetArgs()
+	types := flags.GetArg(1)
 	switch types {
 	case "server":
 		server.Serve()

internal/utils.go

@@ -0,0 +1,101 @@
+package connection
+
+import (
+	"bufio"
+	"bytes"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"log"
+	"os"
+
+	"github.com/keyvchan/NetAssist/pkg/flags"
+	"github.com/keyvchan/NetAssist/pkg/message"
+)
+
+// File represents a abstraced file connection.
+type File struct {
+	Path string
+	File *os.File
+}
+
+// pre-initialize the file connection
+var Stdin = File{
+	Path: "/dev/stdin",
+	File: os.Stdin,
+}
+
+var Stdout = File{
+	Path: "/dev/stdout",
+	File: os.Stdout,
+}
+
+// implements ReadMessage
+func (f File) ReadMessage() message.Message {
+	switch f.Path {
+	case "/dev/stdin":
+		return ReadStdin(f.File)
+
+	default:
+		return message.Message{}
+	}
+}
+
+func ReadStdin(stdin *os.File) message.Message {
+	input_binary := flags.GetArg(4)
+	scanner := bufio.NewScanner(stdin)
+	if scanner.Scan() {
+
+		buf := []byte{}
+		if input_binary == "--binary" {
+			byte_slices := bytes.Split(scanner.Bytes(), []byte(" "))
+			for _, byte_slice := range byte_slices {
+				new_byte := make([]byte, 1024)
+				n, err := hex.Decode(new_byte, byte_slice)
+				if err != nil {
+					// hex parse error, ignore this byte
+					log.Println(err)
+					continue
+				}
+				buf = append(buf, new_byte[:n]...)
+			}
+
+		} else {
+			buf = scanner.Bytes()
+		}
+		// split by space
+
+		return message.Message{
+			Content: buf,
+			Addr:    nil,
+		}
+	} else {
+		log.Fatal(errors.New("failed to read from stdin"))
+	}
+	return message.Message{}
+}
+
+func (f File) WriteMessage(msg message.Message) {
+	switch f.Path {
+	case "/dev/stdout":
+		WriteStdout(f.File, msg)
+	default:
+		return
+	}
+}
+
+func WriteStdout(writter interface{}, message message.Message) {
+
+	// write to stdout
+	input_binary := flags.GetArg(4)
+	fmt.Println(message.Addr)
+
+	if input_binary == "--binary" {
+		for i := 0; i < len(message.Content); i++ {
+			fmt.Printf("%02x ", message.Content[i])
+		}
+		fmt.Print("\n")
+	} else {
+		fmt.Println(string(message.Content))
+	}
+}

internal/write.go

@@ -0,0 +1,43 @@
+package connection
+
+import (
+	"errors"
+	"log"
+	"net"
+
+	"github.com/keyvchan/NetAssist/pkg/message"
+)
+
+type Packet struct {
+	Type string
+	net.PacketConn
+}
+
+func (s Packet) ReadMessage() message.Message {
+	return ReadPacketConn(s.PacketConn)
+}
+
+func ReadPacketConn(reader net.PacketConn) message.Message {
+	buf := make([]byte, 1024)
+	n, addr, err := reader.ReadFrom(buf)
+	if err != nil {
+		log.Fatal(errors.New("error reading from packet"))
+	}
+	return message.Message{
+		Content: buf[:n],
+		Addr:    addr,
+	}
+
+}
+
+// implements WriteMessage
+func (s Packet) WriteMessage(msg message.Message) {
+	WritePacketConn(s.PacketConn, msg)
+}
+
+func WritePacketConn(conn net.PacketConn, message message.Message) {
+	_, err := conn.WriteTo(message.Content, message.Addr)
+	if err != nil {
+		log.Fatal(err)
+	}
+}

main.go

@@ -0,0 +1,60 @@
+package connection
+
+import (
+	"errors"
+	"io"
+	"log"
+	"net"
+
+	"github.com/keyvchan/NetAssist/pkg/message"
+)
+
+type Stream struct {
+	Type string
+	net.Conn
+}
+
+// implements ReadMessage for both type
+func (s Stream) ReadMessage() message.Message {
+	return ReadConn(s.Conn)
+}
+
+var ClosedConn = new(chan net.Conn)
+
+func ReadConn(conn interface{}) message.Message {
+	// type checking
+	connn, ok := conn.(net.Conn)
+	if !ok {
+		log.Fatal("Wrong type")
+	}
+	buf := make([]byte, 1024)
+	// input_binary := GetArg(4)
+	n, err := connn.Read(buf)
+	if errors.Is(err, io.EOF) {
+		log.Println("Connection closed")
+		// remove from channel
+		*ClosedConn <- connn
+		return message.Message{}
+	}
+	if err != nil {
+		log.Println(err)
+		return message.Message{}
+	}
+	message := message.Message{
+		Addr:    connn.RemoteAddr(),
+		Content: buf[:n],
+	}
+	return message
+}
+
+// implements WriteMessage
+func (s Stream) WriteMessage(msg message.Message) {
+	WriteConn(s.Conn, msg)
+}
+
+func WriteConn(conn net.Conn, message message.Message) {
+	_, err := conn.Write(message.Content)
+	if err != nil {
+		log.Fatal(err)
+	}
+}