Fixed several Bugs and added Bitmap for LostPackets

Makefile

@@ -1,5 +1,5 @@
 build:
-	go build -o bin/helloGo main.go client.go packets.go server.go
+	go build -o bin/uftp main.go
 
 server:
 	go run main.go server 

internal/client/client.go

@@ -7,8 +7,9 @@ import (
 	"fmt"
 	"net"
 	"os"
-	"sort"
 	"time"
+
+	"github.com/kelindar/bitmap"
 )
 
 func SendPacket(pck *common.Packet, key [32]byte, conn *net.UDPConn) {
@@ -19,7 +20,7 @@ func SendPacket(pck *common.Packet, key [32]byte, conn *net.UDPConn) {
 }
 
 func ReceivePacket(key [32]byte, conn *net.UDPConn) common.Packet {
-	bytes := make([]byte, common.PacketSize+common.SecureHeaderSize)
+	bytes := make([]byte, common.PacketSize)
 	_, _, err := conn.ReadFrom(bytes)
 	if err != nil {
 		panic(err)
@@ -29,6 +30,7 @@ func ReceivePacket(key [32]byte, conn *net.UDPConn) common.Packet {
 	pck, err := secPck.ExtractPacket(key)
 
 	if err != nil {
+		fmt.Println(bytes)
 		panic(err)
 	}
 
@@ -114,9 +116,10 @@ func GetFile(path string) {
 	ackPck := common.NewAck(&pck)
 	SendPacket(ackPck, key, conn)
 
-	recvPackets := make([]uint32, 0)
 	var endPacket common.Packet
 
+	var recvPackets bitmap.Bitmap
+
 	for {
 		pck := ReceivePacket(key, conn)
 		if pck.Flag == common.End {
@@ -128,10 +131,9 @@ func GetFile(path string) {
 			continue
 		}
 
-		recvPackets = append(recvPackets, pck.Sync)
+		recvPackets.Set(pck.Sync)
 
-		offset := (int64(pck.Sync) - int64(ackPck.Sync+1)) * (common.PacketSize - int64(common.HeaderSize))
+		offset := (int64(pck.Sync) - int64(ackPck.Sync+1)) * int64(common.MaxDataSize)
-		// fmt.Printf("Sync: %v, Offset: %v\n", pck.Sync, offset)
 
 		_, err = file.WriteAt(pck.Data, offset)
 		if err != nil {
@@ -139,19 +141,23 @@ func GetFile(path string) {
 		}
 	}
 
-	sort.Slice(recvPackets, func(i, j int) bool {
-		pckI := recvPackets[i]
-		pckJ := recvPackets[j]
-		return pckI < pckJ
-	})
-
 	lostPackets := make([]uint32, 0)
 
-	for i := ackPck.Sync + 1; i < endPacket.Sync; i++ {
+	var reverse bitmap.Bitmap
-		if b, _ := contains(recvPackets, i); !b {
+	reverse.Grow(endPacket.Sync)
-			lostPackets = append(lostPackets, i)
+	reverse.Ones()
+
+	for i := 0; i <= int(ackPck.Sync); i++ {
+		recvPackets.Set(uint32(i))
 	}
+
+	recvPackets.Xor(reverse)
+
+	recvPackets.Range(func(x uint32) {
+		if x < endPacket.Sync {
+			lostPackets = append(lostPackets, x)
 		}
+	})
 
 	for _, i := range lostPackets {
 		fmt.Println(i)
@@ -177,7 +183,7 @@ func GetFile(path string) {
 				continue
 			}
 
-			offset := (int64(pck.Sync) - int64(ackPck.Sync+1)) * (common.PacketSize - int64(common.HeaderSize))
+			offset := (int64(pck.Sync) - int64(ackPck.Sync+1)) * int64(common.MaxDataSize)
 			// fmt.Printf("Sync: %v, Offset: %v\n", pck.sync, offset)
 
 			_, err = file.WriteAt(pck.Data, offset)

internal/client/go.mod

@@ -1,3 +1,11 @@
 module client
 
 go 1.21.1
+
+require github.com/kelindar/bitmap v1.5.2
+
+require (
+	github.com/kelindar/simd v1.1.2 // indirect
+	github.com/klauspost/cpuid/v2 v2.2.4 // indirect
+	golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e // indirect
+)

internal/client/go.sum

@@ -0,0 +1,8 @@
+github.com/kelindar/bitmap v1.5.2 h1:XwX7CTvJtetQZ64zrOkApoZZHBJRkjE23NfqUALA/HE=
+github.com/kelindar/bitmap v1.5.2/go.mod h1:j3qZjxH9s4OtvsnFTP2bmPkjqil9Y2xQlxPYHexasEA=
+github.com/kelindar/simd v1.1.2 h1:KduKb+M9cMY2HIH8S/cdJyD+5n5EGgq+Aeeleos55To=
+github.com/kelindar/simd v1.1.2/go.mod h1:inq4DFudC7W8L5fhxoeZflLRNpWSs0GNx6MlWFvuvr0=
+github.com/klauspost/cpuid/v2 v2.2.4 h1:acbojRNwl3o09bUq+yDCtZFc1aiwaAAxtcn8YkZXnvk=
+github.com/klauspost/cpuid/v2 v2.2.4/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
+golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e h1:CsOuNlbOuf0mzxJIefr6Q4uAUetRUwZE4qt7VfzP+xo=
+golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

internal/common/packets.go

@@ -11,10 +11,12 @@ import (
 
 const PacketSize = 504
 
-const HeaderSize int = 32 + 1 + 4 + 4
+const HeaderSize int = 1 + 4
-const SecureHeaderSize int = 1 + 42 + 32 + 4
+const SecureHeaderSize int = 1 + 24 + 8 + 4
 
-type SessionID [32]byte
+const MaxDataSize = PacketSize - HeaderSize - SecureHeaderSize - 16 // AEAD Overhead
+
+type SessionID [8]byte
 
 type SecurePacket struct {
 	IsRsa         byte // 0 = false everything else is true
@@ -25,12 +27,13 @@ type SecurePacket struct {
 }
 
 type Packet struct {
-	// headerLength uint32
-	Sid        SessionID
 	Flag HeaderFlag
 	Sync uint32
-	DataLength uint32
 	Data []byte
+
+	// NOT IN BYTES THAT ARE SENT
+	Sid        SessionID
+	DataLength uint32
 }
 
 func NewSymetricSecurePacket(key [32]byte, pck *Packet) *SecurePacket {
@@ -61,9 +64,9 @@ func NewSymetricSecurePacket(key [32]byte, pck *Packet) *SecurePacket {
 func SecurePacketFromBytes(bytes []byte) SecurePacket {
 	isRsa := bytes[0]
 	nonce := bytes[1:25]
-	sid := SessionID(bytes[25:57])
+	sid := SessionID(bytes[25:33])
-	length := binary.LittleEndian.Uint32(bytes[57:61])
+	length := binary.LittleEndian.Uint32(bytes[33:37])
-	enc := bytes[61 : 61+length]
+	enc := bytes[37 : SecureHeaderSize+int(length)]
 
 	return SecurePacket{
 		IsRsa:         isRsa,
@@ -75,12 +78,14 @@ func SecurePacketFromBytes(bytes []byte) SecurePacket {
 }
 
 func (secPck *SecurePacket) ToBytes() []byte {
-	arr := make([]byte, SecureHeaderSize+len(secPck.EncryptedData))
+	encSize := int(secPck.DataLength)
+
+	arr := make([]byte, SecureHeaderSize+encSize)
 	arr[0] = secPck.IsRsa
 	copy(arr[1:25], secPck.Nonce[:])
-	copy(arr[25:57], secPck.Sid[:])
+	copy(arr[25:33], secPck.Sid[:])
-	binary.LittleEndian.PutUint32(arr[57:61], secPck.DataLength)
+	binary.LittleEndian.PutUint32(arr[33:37], secPck.DataLength)
-	copy(arr[61:], secPck.EncryptedData)
+	copy(arr[37:SecureHeaderSize+encSize], secPck.EncryptedData)
 
 	return arr
 }
@@ -95,7 +100,7 @@ func (secPck *SecurePacket) ExtractPacket(key [32]byte) (Packet, error) {
 		return Packet{}, err
 	}
 	// fmt.Println(data)
-	packet := PacketFromBytes(data)
+	packet := PacketFromBytes(data, secPck.DataLength-uint32(HeaderSize)-uint32(aead.Overhead()), secPck.Sid)
 	return packet, nil
 }
 
@@ -105,6 +110,7 @@ func NewRsaPacket(sid SessionID, key [32]byte) *SecurePacket {
 		Nonce:         [24]byte(make([]byte, 24)),
 		Sid:           sid,
 		EncryptedData: key[:],
+		DataLength:    32,
 	}
 }
 
@@ -112,11 +118,9 @@ func (secPck *SecurePacket) ExtractKey( /*RSA HERE LATER*/ ) []byte {
 	return secPck.EncryptedData[:32]
 }
 
-func PacketFromBytes(bytes []byte) Packet {
+func PacketFromBytes(bytes []byte, dataLength uint32, sid SessionID) Packet {
 	flag := HeaderFlag(bytes[0])
-	sid := SessionID(bytes[1:33])
+	sync := binary.LittleEndian.Uint32(bytes[1:5])
-	sync := binary.LittleEndian.Uint32(bytes[33:37])
-	dataLength := binary.LittleEndian.Uint32(bytes[37:41])
 	pck := Packet{
 		Sid:        sid,
 		Flag:       flag,
@@ -231,10 +235,8 @@ func NewPte(fileSize uint32, lastPck *Packet) *Packet {
 func (pck *Packet) ToBytes() []byte {
 	arr := make([]byte, HeaderSize+int(pck.DataLength))
 	arr[0] = byte(pck.Flag)
-	copy(arr[1:33], pck.Sid[:])
+	binary.LittleEndian.PutUint32(arr[1:5], pck.Sync)
-	binary.LittleEndian.PutUint32(arr[33:37], pck.Sync)
+	copy(arr[HeaderSize:HeaderSize+int(pck.DataLength)], pck.Data)
-	binary.LittleEndian.PutUint32(arr[37:41], pck.DataLength)
-	copy(arr[41:], pck.Data)
 
 	return arr
 }

internal/server/server.go

@@ -75,8 +75,8 @@ func (server *Server) resend(conn *net.UDPConn, addr *net.UDPAddr, pck *common.P
 	defer file.Close()
 
 	// This should be different
-	offset := (int64(resend) - 3) * (common.PacketSize - int64(common.HeaderSize))
+	offset := (int64(resend) - 3) * (int64(common.MaxDataSize))
-	buf := make([]byte, common.PacketSize-common.HeaderSize)
+	buf := make([]byte, common.MaxDataSize)
 
 	_, err = file.ReadAt(buf, offset)
 	if err != nil && !errors.Is(err, io.EOF) {
@@ -138,7 +138,7 @@ func (server *Server) sendData(conn *net.UDPConn, addr *net.UDPAddr, pck *common
 	}
 	defer file.Close()
 
-	buf := make([]byte, common.PacketSize-common.HeaderSize)
+	buf := make([]byte, common.MaxDataSize)
 	filePck := pck
 	for {
 		r, err := file.Read(buf)