Rsa Encryption, probably bugged

.gitignore

@@ -2,3 +2,4 @@ bin/
 testFiles/
 out/
 sha512sums.txt
+pubkey.pem

Makefile

@@ -5,7 +5,7 @@ server:
 	go run cmd/uftp/main.go server 
 
 test:
-	go run cmd/uftp/main.go client testFiles/testFile 0.0.0.0:13374
+	go run cmd/uftp/main.go client testFiles/testFile 0.0.0.0
 
 win:
 	GOOS=windows GOARCH=amd64 go build -o bin/app-amd64.exe cmd/uftp/main.go

internal/client/client.go

@@ -2,19 +2,23 @@ package client
 
 import (
 	"crypto/rand"
+	"crypto/x509"
 	"encoding/hex"
+	"encoding/pem"
+	"errors"
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"time"
 
-	"github.com/Pablu23/Uftp/internal/common"
-
 	"github.com/kelindar/bitmap"
+
+	"github.com/Pablu23/Uftp/internal/common"
 )
 
 func SendPacket(pck *common.Packet, key [32]byte, conn *net.UDPConn) {
-	secPck := common.NewSymetricSecurePacket(key, pck)
+	secPck := common.NewSymmetricSecurePacket(key, pck)
 	if _, err := conn.Write(secPck.ToBytes()); err != nil {
 		panic(err)
 	}
@@ -27,9 +31,11 @@ func ReceivePacket(key [32]byte, conn *net.UDPConn) common.Packet {
 		panic(err)
 	}
 
-	secPck := common.SecurePacketFromBytes(bytes)
+	secPck, err := common.SecurePacketFromBytes(bytes)
+	if err != nil {
+		panic(err)
+	}
 	pck, err := secPck.ExtractPacket(key)
-
 	if err != nil {
 		fmt.Println(bytes)
 		panic(err)
@@ -55,7 +61,10 @@ func ReceivePacketWithTimeout(key [32]byte, conn *net.UDPConn) (common.Packet, b
 		return common.Packet{}, false
 	}
 
-	secPck := common.SecurePacketFromBytes(bytes)
+	secPck, err := common.SecurePacketFromBytes(bytes)
+	if err != nil {
+		panic(err)
+	}
 	pck, err := secPck.ExtractPacket(key)
 	if err != nil {
 		panic(err)
@@ -64,6 +73,33 @@ func ReceivePacketWithTimeout(key [32]byte, conn *net.UDPConn) (common.Packet, b
 	return pck, true
 }
 
+func StartConnection(sid common.SessionID, key [32]byte, address string) {
+	pubkey, err := os.ReadFile("pubkey.pem")
+	if err != nil {
+		panic(err)
+	}
+	block, _ := pem.Decode(pubkey)
+	pKey, _ := x509.ParsePKCS1PublicKey(block.Bytes)
+	keyExchangePck, err := common.NewRsaPacket(pKey, key, sid)
+	if err != nil {
+		panic(err)
+	}
+
+	var d net.Dialer
+	conn, err := d.Dial("tcp", address)
+	if err != nil {
+		panic(err)
+	}
+	defer conn.Close()
+	conn.Write(keyExchangePck.ToBytes())
+
+	var buf [1024]byte
+	_, err = conn.Read(buf[:])
+	if err != nil && !errors.Is(err, io.EOF) {
+		panic(err)
+	}
+}
+
 func GetFile(path string, address string) {
 	request := common.NewRequest(path)
 
@@ -73,11 +109,11 @@ func GetFile(path string, address string) {
 		panic(err)
 	}
 	key := [32]byte(k)
-	keyExchangePck := common.NewRsaPacket(request.Sid, key)
+
+	StartConnection(request.Sid, key, fmt.Sprintf("%v:13375", address))
 
 	// udpAddr, err := net.ResolveUDPAddr("udp", "0.0.0.0:13374")
-	udpAddr, err := net.ResolveUDPAddr("udp", address)
+	udpAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%v:13374", address))
-
 	if err != nil {
 		fmt.Println(err)
 		os.Exit(1)
@@ -90,11 +126,6 @@ func GetFile(path string, address string) {
 		os.Exit(1)
 	}
 
-	_, err = conn.Write(keyExchangePck.ToBytes())
-	if err != nil {
-		panic(err)
-	}
-
 	SendPacket(request, key, conn)
 
 	file, err := os.Create("out/" + hex.EncodeToString(request.Sid[:]) + ".recv")

internal/common/constants.go

@@ -0,0 +1,23 @@
+package common
+
+const PacketSize = 504
+
+const (
+	HeaderSize       int = 1 + 4
+	SecureHeaderSize int = 24 + 8 + 4
+)
+
+const MaxDataSize = PacketSize - HeaderSize - SecureHeaderSize - 16 // AEAD Overhead
+
+type SessionID [8]byte
+
+type HeaderFlag uint8
+
+const (
+	Request HeaderFlag = iota
+	PTE     HeaderFlag = iota
+	Ack     HeaderFlag = iota
+	File    HeaderFlag = iota
+	End     HeaderFlag = iota
+	Resend  HeaderFlag = iota
+)

internal/common/packets.go

@@ -5,27 +5,8 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
-
-	"golang.org/x/crypto/chacha20poly1305"
 )
 
-const PacketSize = 504
-
-const HeaderSize int = 1 + 4
-const SecureHeaderSize int = 1 + 24 + 8 + 4
-
-const MaxDataSize = PacketSize - HeaderSize - SecureHeaderSize - 16 // AEAD Overhead
-
-type SessionID [8]byte
-
-type SecurePacket struct {
-	IsRsa         byte // 0 = false everything else is true
-	Nonce         [24]byte
-	Sid           SessionID
-	DataLength    uint32
-	EncryptedData []byte
-}
-
 type Packet struct {
 	Flag HeaderFlag
 	Sync uint32
@@ -36,88 +17,6 @@ type Packet struct {
 	DataLength uint32
 }
 
-func NewSymetricSecurePacket(key [32]byte, pck *Packet) *SecurePacket {
-	sid := pck.Sid
-	data := pck.ToBytes()
-	aead, err := chacha20poly1305.NewX(key[:])
-	if err != nil {
-		panic(err)
-	}
-
-	nonce := make([]byte, 24)
-	if _, err = rand.Read(nonce); err != nil {
-		panic(err)
-	}
-
-	encrypted := make([]byte, len(data)+aead.Overhead())
-	encrypted = aead.Seal(nil, nonce, data, nil)
-
-	return &SecurePacket{
-		IsRsa:         0,
-		Nonce:         [24]byte(nonce),
-		Sid:           sid,
-		DataLength:    uint32(len(encrypted)),
-		EncryptedData: encrypted,
-	}
-}
-
-func SecurePacketFromBytes(bytes []byte) SecurePacket {
-	isRsa := bytes[0]
-	nonce := bytes[1:25]
-	sid := SessionID(bytes[25:33])
-	length := binary.LittleEndian.Uint32(bytes[33:37])
-	enc := bytes[37 : SecureHeaderSize+int(length)]
-
-	return SecurePacket{
-		IsRsa:         isRsa,
-		Nonce:         [24]byte(nonce),
-		Sid:           sid,
-		DataLength:    length,
-		EncryptedData: enc,
-	}
-}
-
-func (secPck *SecurePacket) ToBytes() []byte {
-	encSize := int(secPck.DataLength)
-
-	arr := make([]byte, SecureHeaderSize+encSize)
-	arr[0] = secPck.IsRsa
-	copy(arr[1:25], secPck.Nonce[:])
-	copy(arr[25:33], secPck.Sid[:])
-	binary.LittleEndian.PutUint32(arr[33:37], secPck.DataLength)
-	copy(arr[37:SecureHeaderSize+encSize], secPck.EncryptedData)
-
-	return arr
-}
-
-func (secPck *SecurePacket) ExtractPacket(key [32]byte) (Packet, error) {
-	aead, err := chacha20poly1305.NewX(key[:])
-	if err != nil {
-		panic(err)
-	}
-	data, err := aead.Open(nil, secPck.Nonce[:], secPck.EncryptedData, nil)
-	if err != nil {
-		return Packet{}, err
-	}
-	// fmt.Println(data)
-	packet := PacketFromBytes(data, secPck.DataLength-uint32(HeaderSize)-uint32(aead.Overhead()), secPck.Sid)
-	return packet, nil
-}
-
-func NewRsaPacket(sid SessionID, key [32]byte) *SecurePacket {
-	return &SecurePacket{
-		IsRsa:         1,
-		Nonce:         [24]byte(make([]byte, 24)),
-		Sid:           sid,
-		EncryptedData: key[:],
-		DataLength:    32,
-	}
-}
-
-func (secPck *SecurePacket) ExtractKey( /*RSA HERE LATER*/ ) []byte {
-	return secPck.EncryptedData[:32]
-}
-
 func PacketFromBytes(bytes []byte, dataLength uint32, sid SessionID) Packet {
 	flag := HeaderFlag(bytes[0])
 	sync := binary.LittleEndian.Uint32(bytes[1:5])
@@ -240,14 +139,3 @@ func (pck *Packet) ToBytes() []byte {
 
 	return arr
 }
-
-type HeaderFlag uint8
-
-const (
-	Request HeaderFlag = iota
-	PTE     HeaderFlag = iota
-	Ack     HeaderFlag = iota
-	File    HeaderFlag = iota
-	End     HeaderFlag = iota
-	Resend  HeaderFlag = iota
-)

internal/common/packets_test.go

@@ -55,9 +55,42 @@ func TestSymetricSecurePacket(t *testing.T) {
 		DataLength: 3,
 	}
 
-	key := [32]byte{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1}
+	key := [32]byte{
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+		1,
+	}
 
-	secPck := NewSymetricSecurePacket(key, &expect)
+	secPck := NewSymmetricSecurePacket(key, &expect)
 
 	packet, err := secPck.ExtractPacket(key)
 	if err != nil {
@@ -71,8 +104,6 @@ func TestSymetricSecurePacket(t *testing.T) {
 
 func TestSecurePacketFromBytes(t *testing.T) {
 	bytes := []byte{
-		//IsRsa
-		0,
 		// Nonce
 		1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
 		// Sid
@@ -83,16 +114,44 @@ func TestSecurePacketFromBytes(t *testing.T) {
 		101, 10, 1,
 	}
 
-	secPck := SecurePacketFromBytes(bytes)
-
 	expect := SecurePacket{
-		IsRsa:         0,
+		Nonce: [24]byte{
-		Nonce:         [24]byte{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+		},
 		Sid:           [8]byte{255, 255, 255, 255, 255, 255, 255, 255},
 		DataLength:    3,
 		EncryptedData: []byte{101, 10, 1},
 	}
 
+	secPck, err := SecurePacketFromBytes(bytes)
+
+	if err != nil {
+		t.Fail()
+	}
+
 	if !cmp.Equal(secPck, expect) {
 		t.Fail()
 	}
@@ -100,8 +159,6 @@ func TestSecurePacketFromBytes(t *testing.T) {
 
 func TestSecurePacketToBytes(t *testing.T) {
 	expect := []byte{
-		//IsRsa
-		0,
 		// Nonce
 		1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
 		// Sid
@@ -113,8 +170,32 @@ func TestSecurePacketToBytes(t *testing.T) {
 	}
 
 	secPck := SecurePacket{
-		IsRsa:         0,
+		Nonce: [24]byte{
-		Nonce:         [24]byte{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+			1,
+		},
 		Sid:           [8]byte{255, 255, 255, 255, 255, 255, 255, 255},
 		DataLength:    3,
 		EncryptedData: []byte{101, 10, 1},

internal/common/rsapacket.go

@@ -0,0 +1,63 @@
+package common
+
+import (
+	"crypto"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/sha256"
+	"encoding/binary"
+)
+
+type RsaPacket struct {
+	Sid          SessionID
+	DataLength   uint32
+	EncryptedKey []byte
+}
+
+func NewRsaPacket(pubKey *rsa.PublicKey, key [32]byte, sid SessionID) (*RsaPacket, error) {
+	enc, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, pubKey, key[:], nil)
+	if err != nil {
+		return nil, err
+	}
+
+	pck := RsaPacket{
+		Sid:          sid,
+		DataLength:   uint32(len(enc)),
+		EncryptedKey: enc,
+	}
+
+	return &pck, nil
+}
+
+func (rsaPck *RsaPacket) ToBytes() []byte {
+	bytes := make([]byte, rsaPck.DataLength+8+4)
+	copy(bytes[0:8], rsaPck.Sid[:])
+	binary.LittleEndian.PutUint32(bytes[8:12], rsaPck.DataLength)
+	copy(bytes[12:], rsaPck.EncryptedKey[:])
+	return bytes
+}
+
+func RsaPacketFromBytes(bytes []byte) *RsaPacket {
+	sid := SessionID(bytes[0:8])
+	dLen := binary.LittleEndian.Uint32(bytes[8:12])
+	data := bytes[12 : 12+dLen]
+
+	return &RsaPacket{
+		Sid:          sid,
+		DataLength:   dLen,
+		EncryptedKey: data,
+	}
+}
+
+func (rsaPck *RsaPacket) ExtractKey(priv *rsa.PrivateKey) ([32]byte, error) {
+	// key, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, priv, rsaPck.EncryptedKey, nil)
+	key, err := priv.Decrypt(
+		rand.Reader,
+		rsaPck.EncryptedKey,
+		&rsa.OAEPOptions{Hash: crypto.SHA256},
+	)
+	if err != nil {
+		return [32]byte{}, err
+	}
+	return [32]byte(key[0:32]), nil
+}

internal/common/securepacket.go

@@ -0,0 +1,87 @@
+package common
+
+import (
+	"crypto/rand"
+	"encoding/binary"
+	"errors"
+
+	"golang.org/x/crypto/chacha20poly1305"
+)
+
+type SecurePacket struct {
+	Nonce         [24]byte
+	Sid           SessionID
+	DataLength    uint32
+	EncryptedData []byte
+}
+
+func NewSymmetricSecurePacket(key [32]byte, pck *Packet) *SecurePacket {
+	sid := pck.Sid
+	data := pck.ToBytes()
+	aead, err := chacha20poly1305.NewX(key[:])
+	if err != nil {
+		panic(err)
+	}
+
+	nonce := make([]byte, 24)
+	if _, err = rand.Read(nonce); err != nil {
+		panic(err)
+	}
+
+	encrypted := make([]byte, len(data)+aead.Overhead())
+	encrypted = aead.Seal(nil, nonce, data, nil)
+
+	return &SecurePacket{
+		Nonce:         [24]byte(nonce),
+		Sid:           sid,
+		DataLength:    uint32(len(encrypted)),
+		EncryptedData: encrypted,
+	}
+}
+
+func SecurePacketFromBytes(bytes []byte) (*SecurePacket, error) {
+	nonce := bytes[:24]
+	sid := SessionID(bytes[24:32])
+	length := binary.LittleEndian.Uint32(bytes[32:36])
+	if SecureHeaderSize+int(length) > PacketSize {
+		return nil, errors.New("Packet too large")
+	}
+	enc := bytes[36 : SecureHeaderSize+int(length)]
+
+	return &SecurePacket{
+		Nonce:         [24]byte(nonce),
+		Sid:           sid,
+		DataLength:    length,
+		EncryptedData: enc,
+	}, nil
+}
+
+func (secPck *SecurePacket) ToBytes() []byte {
+	encSize := int(secPck.DataLength)
+
+	arr := make([]byte, SecureHeaderSize+encSize)
+	copy(arr[0:24], secPck.Nonce[:])
+	copy(arr[24:32], secPck.Sid[:])
+	binary.LittleEndian.PutUint32(arr[32:36], secPck.DataLength)
+	copy(arr[36:SecureHeaderSize+encSize], secPck.EncryptedData)
+
+	return arr
+}
+
+func (secPck *SecurePacket) ExtractPacket(key [32]byte) (Packet, error) {
+	aead, err := chacha20poly1305.NewX(key[:])
+	if err != nil {
+		panic(err)
+	}
+	data, err := aead.Open(nil, secPck.Nonce[:], secPck.EncryptedData, nil)
+	if err != nil {
+		return Packet{}, err
+	}
+	// fmt.Println(data)
+	packet := PacketFromBytes(
+		data,
+		secPck.DataLength-uint32(HeaderSize)-uint32(aead.Overhead()),
+		secPck.Sid,
+	)
+	return packet, nil
+}

internal/server/server.go

@@ -1,9 +1,12 @@
 package server
 
 import (
+	"bufio"
 	"crypto/rand"
 	"crypto/rsa"
+	"crypto/x509"
 	"encoding/hex"
+	"encoding/pem"
 	"errors"
 	"io"
 	"net"
@@ -12,9 +15,9 @@ import (
 	"sync"
 	"time"
 
-	"github.com/Pablu23/Uftp/internal/common"
-
 	log "github.com/sirupsen/logrus"
+
+	"github.com/Pablu23/Uftp/internal/common"
 )
 
 type info struct {
@@ -33,7 +36,6 @@ type Server struct {
 
 func New() (*Server, error) {
 	key, err := rsa.GenerateKey(rand.Reader, 4096)
-
 	if err != nil {
 		return nil, err
 	}
@@ -48,6 +50,20 @@ func New() (*Server, error) {
 	}, nil
 }
 
+func (server *Server) SavePublicKeyPem() error {
+	file, err := os.Create("pubkey.pem")
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+	publicKeyPEM := &pem.Block{
+		Type:  "RSA PUBLIC KEY",
+		Bytes: x509.MarshalPKCS1PublicKey(&server.rsa.PublicKey),
+	}
+	pem.Encode(file, publicKeyPEM)
+	return nil
+}
+
 func (server *Server) sendPacket(conn *net.UDPConn, addr *net.UDPAddr, pck *common.Packet) {
 	server.mu.Lock()
 	var key [32]byte
@@ -61,7 +77,7 @@ func (server *Server) sendPacket(conn *net.UDPConn, addr *net.UDPAddr, pck *comm
 		return
 	}
 
-	secPck := common.NewSymetricSecurePacket(key, pck)
+	secPck := common.NewSymmetricSecurePacket(key, pck)
 	if _, err := conn.WriteToUDP(secPck.ToBytes(), addr); err != nil {
 		log.Error("Could not write Packet to UDP")
 		return
@@ -154,7 +170,6 @@ func (server *Server) handleAck(conn *net.UDPConn, addr *net.UDPAddr, pck *commo
 		server.mu.Unlock()
 		return
 	}
-
 }
 
 func (server *Server) sendPTE(conn *net.UDPConn, addr *net.UDPAddr, pck *common.Packet) {
@@ -189,7 +204,6 @@ func (server *Server) sendPTE(conn *net.UDPConn, addr *net.UDPAddr, pck *common.
 }
 
 func (server *Server) sendData(conn *net.UDPConn, addr *net.UDPAddr, pck *common.Packet) {
-
 	var path string
 	server.mu.Lock()
 	if info, ok := server.sessions[pck.Sid]; ok {
@@ -282,9 +296,51 @@ func (server *Server) handleShutdown(stop chan bool) {
 	}()
 }
 
+func (server *Server) handleConnection(conn net.Conn) {
+	reader := bufio.NewReader(conn)
+	var buf [2048]byte
+	r, err := reader.Read(buf[:])
+	if err != nil {
+		log.WithError(err).Warn("Could not read from Connection")
+		conn.Close()
+		return
+	}
+
+	// fmt.Println(buf)
+
+	rsaPck := common.RsaPacketFromBytes(buf[0:r])
+	key, err := rsaPck.ExtractKey(server.rsa)
+	if err != nil && !errors.Is(err, io.EOF) {
+		log.WithError(err).Warn("Could not extract Key")
+		return
+	}
+	server.mu.Lock()
+	server.sessions[rsaPck.Sid] = &info{
+		key: key,
+	}
+	server.mu.Unlock()
+	conn.Write([]byte("Yep"))
+	conn.Close()
+}
+
+func (server *Server) startManagement() {
+	listener, err := net.Listen("tcp", "0.0.0.0:13375")
+	if err != nil {
+		log.Fatal("Could not start listening on TCP 0.0.0.0:13375")
+	}
+	defer listener.Close()
+	for {
+		conn, err := listener.Accept()
+		if err != nil {
+			log.WithError(err).Warn("Could not accept TCP Connection")
+		}
+
+		go server.handleConnection(conn)
+	}
+}
+
 func (server *Server) Serve() {
 	udpAddr, err := net.ResolveUDPAddr("udp", "0.0.0.0:13374")
-
 	if err != nil {
 		log.Fatal("Could not resolve UDP Address")
 	}
@@ -301,6 +357,7 @@ func (server *Server) Serve() {
 	c := make(chan bool)
 	server.handleShutdown(c)
 	go server.startTimeout(c)
+	go server.startManagement()
 
 	for {
 		var buf [common.PacketSize]byte
@@ -310,11 +367,13 @@ func (server *Server) Serve() {
 			continue
 		}
 
-		secPck := common.SecurePacketFromBytes(buf[:])
+		secPck, err := common.SecurePacketFromBytes(buf[:])
+		if err != nil {
+			log.WithError(err).Warn("Received invalid Packet")
+			continue
+		}
 
-		if secPck.IsRsa == 0 {
 		var key [32]byte
-
 		server.mu.Lock()
 		if info, ok := server.sessions[secPck.Sid]; ok {
 			key = info.key
@@ -323,19 +382,12 @@ func (server *Server) Serve() {
 			server.mu.Unlock()
 			continue
 		}
-			server.mu.Unlock()
 		pck, err := secPck.ExtractPacket(key)
 		if err != nil {
 			log.Error("Could not extract Packet from Secure Packet")
 		}
+		server.mu.Unlock()
 		go server.handlePacket(conn, addr, &pck)
-		} else {
-			key := secPck.ExtractKey()
-			log.WithField("SessionID", hex.EncodeToString(secPck.Sid[:])).Info("New Session")
-			server.sessions[secPck.Sid] = &info{
-				key: [32]byte(key),
-			}
-		}
 
 	}
 }