working on milestone 2

4 files changed, 296 insertions, 15 deletions

diff --git a/cmd/vhost/main.go b/cmd/vhost/main.go
index 57f5f7c..7c32df8 100644
--- a/cmd/vhost/main.go
+++ b/cmd/vhost/main.go
@@ -25,6 +25,10 @@ func main() {
 	}
 	ipstack.RegisterProtocolHandler(ipstack.TEST_PROTOCOL)
 	ipstack.RegisterProtocolHandler(ipstack.TCP_PROTOCOL)
+
+	// create a map of sockets
+	sockets := make(map[int]*ipstack.VTCPConn)
+
 	scanner := bufio.NewScanner(os.Stdin)
 
 	for scanner.Scan() {
@@ -97,12 +101,20 @@ func main() {
 				fmt.Println(err)
 				continue
 			}
-			_, err = listener.VAccept()
-			if err != nil {
-				fmt.Println(err)
+			go func() {
+				for {
+					conn, err := listener.VAccept()
+					if err != nil {
+						fmt.Println(err)
+					}
+					sockets[conn.Socket] = conn
+				}
+			}() 
+		case "c":
+			if len(tokens) != 3 {
+				fmt.Println("Invalid command: ", line)
 				continue
 			}
-		case "c":
 			// connect to a port
 			vipAddr := tokens[1]
 			port := tokens[2]
@@ -111,15 +123,76 @@ func main() {
 				fmt.Println(err)
 				continue
 			}
-			_, err = ipstack.VConnect(vipAddr, uint16(uint16Port))
+			conn, err := ipstack.VConnect(vipAddr, uint16(uint16Port))
 			if err != nil {
 				fmt.Println(err)
 				continue
 			}
+			sockets[conn.Socket] = conn
 		case "ls":
 			// list sockets
 			fmt.Println("SID\tLAddr\t\tLPort\tRAddr\t\tRPort\tRStatus")
 			fmt.Println(ipstack.SprintSockets())
+		case "s":
+			if len(tokens) < 3 {
+				fmt.Println("Invalid command: ", line)
+				continue
+			}
+			socketID, err := strconv.Atoi(tokens[1])
+			if err != nil {
+				fmt.Println(err)
+				continue
+			}
+			conn, _ := sockets[socketID]
+			message := tokens[2:]
+			messageToSend := strings.Join(message, " ")
+			messageToSendBytes := []byte(messageToSend)
+			bytesWritten, err := conn.VWrite(messageToSendBytes)
+			if err != nil {
+				fmt.Println(err)
+				continue
+			}
+			fmt.Printf("Sent %d bytes\n", bytesWritten)
+		case "r":
+			if len(tokens) != 3 {
+				fmt.Println("Invalid command: ", line)
+				continue
+			}
+			socketID, err := strconv.Atoi(tokens[1])
+			if err != nil {
+				fmt.Println(err)
+				continue
+			}
+			numberOfBytes, err := strconv.Atoi(tokens[2])
+			if err != nil {
+				fmt.Println(err)
+				continue
+			}
+			conn, _ := sockets[socketID]
+			// buffer := make([]byte, numberOfBytes)
+			bytesRead, message, err := conn.VRead(numberOfBytes)
+			if err != nil {
+				fmt.Println(err)
+				continue
+			}
+			fmt.Printf("Read %d bytes: %s\n", bytesRead, message)
+		case "cl":
+			// close a socket
+			if len(tokens) != 2 {
+				fmt.Println("Invalid command: ", line)
+				continue
+			}
+			socketID, err := strconv.Atoi(tokens[1])
+			if err != nil {
+				fmt.Println(err)
+				continue
+			}
+			conn, _ := sockets[socketID]
+			err = conn.VClose()
+			if err != nil {
+				fmt.Println(err)
+				continue
+			}
 		default:
 			fmt.Println("Invalid command: ", line)
 			fmt.Println("Commands: ")
diff --git a/pkg/ipstack/ipstack.go b/pkg/ipstack/ipstack.go
index ec87485..0d317c2 100644
--- a/pkg/ipstack/ipstack.go
+++ b/pkg/ipstack/ipstack.go
@@ -891,9 +891,37 @@ func HandleTCP(src *Interface, message []byte, hdr *ipv4header.IPv4Header) error
 		// if the SYN flag is set, then send a SYNACK
 		available := false
 
+		socketEntry, in := VHostSocketMaps[SocketKey{hdr.Dst.String(), tcpHdr.DstPort, hdr.Src.String(), tcpHdr.SrcPort}]
+		if !in {
+			fmt.Println("no socket entry found")
+		} else if socketEntry.State == Established {
+			fmt.Println("socket entry found")
+			
+			// make ack header
+			tcpHdr := &header.TCPFields{
+				SrcPort: tcpHdr.DstPort,
+				DstPort: tcpHdr.SrcPort,
+				SeqNum:  tcpHdr.SeqNum,
+				AckNum:  tcpHdr.SeqNum + 1,
+				DataOffset: 20,
+				Flags:   0x10,
+				WindowSize: MAX_WINDOW_SIZE,
+				Checksum: 0,
+				UrgentPointer: 0,
+			}
+			// make the payload
+			err := SendTCP(tcpHdr, message, hdr.Dst, hdr.Src)
+			if err != nil {
+				fmt.Println(err)
+			}
+			socketEntry.Conn.RecvBuffer.buffer = append(socketEntry.Conn.RecvBuffer.buffer, tcpPayload...)
+			socketEntry.Conn.RecvBuffer.recvNext += uint32(len(tcpPayload))
+			break
+		}
 		// add to table if available
 		mapMutex.Lock()
 		for _, socketEntry := range VHostSocketMaps {
+			// todo: check between all 4 field in tuple
 			if socketEntry.LocalPort == tcpHdr.DstPort && socketEntry.LocalIP == hdr.Dst.String() && socketEntry.State == Listening{
 				// add a new socketEntry to the map
 				newEntry := &SocketEntry{
@@ -904,7 +932,9 @@ func HandleTCP(src *Interface, message []byte, hdr *ipv4header.IPv4Header) error
 					State: SYNRECIEVED,
 					Socket: socketsMade,
 				}
-				VHostSocketMaps[socketsMade] = newEntry
+				// add the entry to the map
+				key := SocketKey{hdr.Dst.String(), tcpHdr.DstPort, hdr.Src.String(), tcpHdr.SrcPort}
+				VHostSocketMaps[key] = newEntry
 				socketsMade += 1
 				// add the entry to the map
 				available = true
@@ -954,8 +984,8 @@ func HandleTCP(src *Interface, message []byte, hdr *ipv4header.IPv4Header) error
 		tcpHdr := &header.TCPFields{
 			SrcPort: tcpHdr.DstPort,
 			DstPort: tcpHdr.SrcPort,
-			SeqNum:  tcpHdr.SeqNum,
-			AckNum:  tcpHdr.SeqNum + 1,
+			SeqNum:  tcpHdr.SeqNum + 1,
+			AckNum:  tcpHdr.SeqNum,
 			DataOffset: 20,
 			Flags:   0x10,
 			WindowSize: MAX_WINDOW_SIZE,
@@ -973,6 +1003,16 @@ func HandleTCP(src *Interface, message []byte, hdr *ipv4header.IPv4Header) error
 		fmt.Println("I see an ACK flag")
 		// lookup for socket entry and update its state
 		// set synChan to true (TODO)
+		key := SocketKey{hdr.Dst.String(), tcpHdr.DstPort, hdr.Src.String(), tcpHdr.SrcPort}
+		socketEntry, in := VHostSocketMaps[key]
+		if !in {
+			fmt.Println("no socket entry found")
+		} else if (socketEntry.State == Established) {
+			fmt.Println("socket entry found")
+			// socketEntry.Conn.RecvBuffer.buffer = append(socketEntry.Conn.RecvBuffer.buffer, tcpPayload...)
+			socketEntry.Conn.SendBuffer.una += uint32(len(tcpPayload))
+			break
+		}
 		
 		mapMutex.Lock()
 		for _, socketEntry := range VHostSocketMaps {
@@ -1176,7 +1216,8 @@ type VTCPConn struct {
 	RemotePort uint16
 	Socket int
 	State ConnectionState
-	Buffer []byte
+	SendBuffer *SendBuffer
+	RecvBuffer *RecvBuffer
 }
 
 type SocketEntry struct {
@@ -1186,10 +1227,32 @@ type SocketEntry struct {
 	RemoteIP  string
 	RemotePort uint16
 	State     ConnectionState
+	Conn 	*VTCPConn
+}
+
+type SocketKey struct {
+	LocalIP   string
+	LocalPort uint16
+	RemoteIP  string
+	RemotePort uint16
+}
+
+type RecvBuffer struct {
+	recvNext uint32
+	lbr uint32
+	buffer []byte
+}
+
+type SendBuffer struct {
+	una uint32
+	nxt uint32
+	lbr uint32
+	buffer []byte
 }
 
 // create a socket map 
-var VHostSocketMaps = make(map[int]*SocketEntry)
+// var VHostSocketMaps = make(map[int]*SocketEntry)
+var VHostSocketMaps = make(map[SocketKey]*SocketEntry)
 // create a channel map 
 var VHostChannelMaps = make(map[int]chan []byte)
 var mapMutex = &sync.Mutex{}
@@ -1205,10 +1268,12 @@ func VListen(port uint16) (*VTCPListener, error) {
 		LocalPort: port,
 		LocalAddr: myIP.String(),
 	}
-	
+
 	// add the socket to the socket map
 	mapMutex.Lock()
-	VHostSocketMaps[socketsMade] = &SocketEntry{
+
+	key := SocketKey{myIP.String(), port, "", 0}
+	VHostSocketMaps[key] = &SocketEntry{
 		Socket: socketsMade,
 		LocalIP:  myIP.String(),
 		LocalPort: port,
@@ -1237,7 +1302,19 @@ func (l *VTCPListener) VAccept() (*VTCPConn, error) {
 					RemotePort: socketEntry.RemotePort,
 					Socket: socketEntry.Socket,
 					State: Established,
+					SendBuffer: &SendBuffer{
+						una: 0,
+						nxt: 0,
+						lbr: 0,
+						buffer: make([]byte, MAX_WINDOW_SIZE),
+					},
+					RecvBuffer: &RecvBuffer{
+						recvNext: 0,
+						lbr: 0,
+						buffer: make([]byte, MAX_WINDOW_SIZE),
+					},
 				}
+				socketEntry.Conn = conn
 				mapMutex.Unlock()
 				return conn, nil
 			}
@@ -1265,7 +1342,7 @@ func VConnect(ip string, port uint16) (*VTCPConn, error) {
 		SrcPort: portRand,
 		DstPort: port,
 		SeqNum:  startingSeqNum,
-		AckNum:  startingSeqNum,
+		AckNum:  0,
 		DataOffset: 20,
 		Flags:   header.TCPFlagSyn,
 		WindowSize: MAX_WINDOW_SIZE,
@@ -1277,6 +1354,7 @@ func VConnect(ip string, port uint16) (*VTCPConn, error) {
 	if err != nil {
 		return nil, err
 	}
+
 	err = SendTCP(tcpHdr, payload, myIP, ipParsed)
 	if err != nil {
 		return nil, err
@@ -1289,18 +1367,30 @@ func VConnect(ip string, port uint16) (*VTCPConn, error) {
 		RemotePort: port,
 		Socket: socketsMade,
 		State: Established,
-		Buffer: []byte{},
+		SendBuffer: &SendBuffer{
+			una: 0,
+			nxt: 0,
+			lbr: 0,
+			buffer: make([]byte, MAX_WINDOW_SIZE),
+		},
+		RecvBuffer: &RecvBuffer{
+			recvNext: 0,
+			lbr: 0,
+			buffer: make([]byte, MAX_WINDOW_SIZE),
+		},
 	}
 
 	// add the socket to the socket map
+	key := SocketKey{myIP.String(), portRand, ip, port}
 	mapMutex.Lock()
-	VHostSocketMaps[socketsMade] = &SocketEntry{
+	VHostSocketMaps[key] = &SocketEntry{
 		Socket: socketsMade,
 		LocalIP:  myIP.String(),
 		LocalPort: portRand,
 		RemoteIP:  ip,
 		RemotePort: port,
 		State:     SYNSENT,
+		Conn: conn,
 	}
 	mapMutex.Unlock()
 	socketsMade += 1
@@ -1354,4 +1444,122 @@ func SprintSockets() string {
 		tmp += fmt.Sprintf("%d\t%s\t%d\t%s\t%d\t%s\n", socket.Socket, socket.LocalIP, socket.LocalPort, socket.RemoteIP, socket.RemotePort, socket.State)
 	}
 	return tmp
+}
+
+// MILESTONE 2
+func (c *VTCPConn) VClose() error {
+	// check if the socket is in the map
+	key := SocketKey{c.LocalAddr, c.LocalPort, c.RemoteAddr, c.RemotePort}
+	mapMutex.Lock()
+	socketEntry, in := VHostSocketMaps[key]
+	mapMutex.Unlock()
+	if !in {
+		return errors.Errorf("error VClose: socket %d does not exist", c.Socket)
+	}
+
+	// change the state to closed
+	socketEntry.State = Closed
+	return nil
+}
+
+
+// advertise window = max window size - (next - 1 - lbr)
+
+// early arrivals queue
+var earlyArrivals = make([][]byte, 0)
+
+// retranmission queue
+var retransmissionQueue = make([][]byte, 0)
+
+func (c *VTCPConn) VWrite(payload []byte) (int, error) {
+	// check if the socket is in the map
+	key := SocketKey{c.LocalAddr, c.LocalPort, c.RemoteAddr, c.RemotePort}
+	mapMutex.Lock()
+	socketEntry, in := VHostSocketMaps[key]
+	mapMutex.Unlock()
+	if !in {
+		return 0, errors.Errorf("error VWrite: socket %d does not exist", c.Socket)
+	}
+
+	// check if the state is established
+	if socketEntry.State != Established {
+		return 0, errors.Errorf("error VWrite: socket %d is not in established state", c.Socket)
+	}
+
+	// check if the payload is empty
+	if len(payload) == 0 {
+		return 0, nil
+	}
+
+	// check if the payload is larger than the window size
+	if len(payload) > MAX_WINDOW_SIZE {
+		return 0, errors.Errorf("error VWrite: payload is larger than the window size")
+	}
+
+	// check if the payload is larger than the available window size
+	if len(payload) > int(MAX_WINDOW_SIZE - (c.SendBuffer.nxt - 1 - c.SendBuffer.lbr)) {
+		return 0, errors.Errorf("error VWrite: payload is larger than the available window size")
+	}
+
+	// make the header
+	advertisedWindow := MAX_WINDOW_SIZE - (c.SendBuffer.nxt - 1 - c.SendBuffer.lbr)
+	tcpHdr := &header.TCPFields{
+		SrcPort: c.LocalPort,
+		DstPort: c.RemotePort,
+		SeqNum:  c.SendBuffer.nxt,
+		AckNum:  c.SendBuffer.una,
+		DataOffset: 20,
+		Flags:   header.TCPFlagSyn,
+		WindowSize: uint16(advertisedWindow),
+		Checksum: 0,
+		UrgentPointer: 0,
+	}
+
+	myIP := GetInterfaces()[0].IpPrefix.Addr()
+	ipParsed, err := netip.ParseAddr(c.RemoteAddr)
+	if err != nil {
+		return 0, err
+	}
+
+	err = SendTCP(tcpHdr, payload, myIP, ipParsed)
+	if err != nil {
+		return 0, err
+	}
+	// update the next sequence number
+	// c.SendBuffer.nxt += uint32(len(payload))
+
+
+	c.SendBuffer.lbr += uint32(len(payload))
+	return len(payload), nil
+}
+
+
+func (c *VTCPConn) VRead(numBytesToRead int) (int, string, error) {
+	// check if the socket is in the map
+	key := SocketKey{c.LocalAddr, c.LocalPort, c.RemoteAddr, c.RemotePort}
+	// mapMutex.Lock()
+	socketEntry, in := VHostSocketMaps[key]
+	// mapMutex.Unlock()
+	// check if the socket is in the map
+	if !in {
+		return 0, "", errors.Errorf("error VRead: socket %d does not exist", c.Socket)
+	}
+
+	// check if the state is established
+	if socketEntry.State != Established {
+		return 0, "", errors.Errorf("error VRead: socket %d is not in established state", c.Socket)
+	}
+	fmt.Println("I am in VRead")
+	fmt.Println("I have", c.RecvBuffer.recvNext - c.RecvBuffer.lbr, "bytes to read")
+	fmt.Println(c.RecvBuffer.recvNext, c.RecvBuffer.lbr)
+	if (c.RecvBuffer.lbr < c.RecvBuffer.recvNext && c.RecvBuffer.recvNext - c.RecvBuffer.lbr >= uint32(numBytesToRead)) {
+		fmt.Println("I have enough data to read")
+		toReturn := string(socketEntry.Conn.RecvBuffer.buffer[c.RecvBuffer.lbr:c.RecvBuffer.lbr+uint32(numBytesToRead)])
+		// update the last byte read
+		c.RecvBuffer.lbr += uint32(numBytesToRead)
+		// return the data
+		return numBytesToRead, toReturn, nil
+	}
+
+	return 0, "", nil
 }
\ No newline at end of file
diff --git a/vhost b/vhost
index 5919be7..2397f1e 100755
--- a/vhost
+++ b/vhost
diff --git a/vrouter b/vrouter
index 7708666..f347322 100755
--- a/vrouter
+++ b/vrouter