bismuthd/server/server.go

package server

import (
	"net"

	core "git.greysoh.dev/imterah/bismuthd/commons"

	"crypto/cipher"
	"crypto/rand"

	"github.com/ProtonMail/gopenpgp/v3/crypto"
	"golang.org/x/crypto/chacha20poly1305"
)

// Called to handle a connnection for Bismuth. The conn argument is the client you'd like to handle
func (bismuth BismuthServer) HandleProxy(conn net.Conn) error {
	serverState := "keyHandshake"

	var clientPublicKey *crypto.Key
	var aead cipher.AEAD

	for {
		if serverState == "keyHandshake" {
			dataModeByteArr := make([]byte, 1)
			_, err := conn.Read(dataModeByteArr)

			dataMode := dataModeByteArr[0]

			if err != nil {
				conn.Close()
				return err
			}

			if dataMode == core.SendPublicKey {
				pubKeyLengthBytes := make([]byte, 3)
				_, err := conn.Read(pubKeyLengthBytes)

				if err != nil {
					conn.Close()
					return err
				}

				pubKeyLength := core.Int24ToInt32(pubKeyLengthBytes)

				publicKey := make([]byte, pubKeyLength)
				_, err = conn.Read(publicKey)

				pubKey := publicKey[:]

				// Attempt to parse the public key
				clientPublicKey, err = crypto.NewKey(pubKey)

				if err != nil {
					conn.Close()
					return err
				}

				// Now, we send our key
				pubKey, err = bismuth.PublicKey.GetPublicKey()

				if err != nil {
					conn.Close()
					return err
				}

				pubKeyLength = uint32(len(pubKey))
				core.Int32ToInt24(pubKeyLengthBytes, pubKeyLength)

				conn.Write([]byte{core.SendPublicKey})
				conn.Write(pubKeyLengthBytes)
				conn.Write([]byte(pubKey))

				serverState = "symmHandshake"
			}
		} else if serverState == "symmHandshake" {
			dataModeByteArr := make([]byte, 1)
			_, err := conn.Read(dataModeByteArr)

			dataMode := dataModeByteArr[0]

			if err != nil {
				conn.Close()
				return err
			}

			if dataMode == core.SwitchToSymmetricKey {
				// TODO: Make this not hard-coded
				symmetricKey := make([]byte, chacha20poly1305.KeySize)

				encryptionData := []byte{
					byte(bismuth.SymmetricEncryptionAlgorithm),
				}

				if _, err := rand.Read(symmetricKey); err != nil {
					conn.Close()
					return err
				}

				encryptionData = append(encryptionData, symmetricKey...)

				encHandle, err := bismuth.pgp.Encryption().Recipient(clientPublicKey).New()

				if err != nil {
					conn.Close()
					return err
				}

				pgpMessage, err := encHandle.Encrypt(encryptionData)

				if err != nil {
					conn.Close()
					return err
				}

				encryptedMessage := pgpMessage.Bytes()
				gpgMessageLengthBytes := make([]byte, 3)

				core.Int32ToInt24(gpgMessageLengthBytes, uint32(len(encryptedMessage)))

				conn.Write([]byte{core.SwitchToSymmetricKey})
				conn.Write(gpgMessageLengthBytes)
				conn.Write(encryptedMessage)

				aead, err = chacha20poly1305.NewX(symmetricKey)

				if err != nil {
					conn.Close()
					return err
				}

				serverState = "APITransmit"
			}
		} else if serverState == "APITransmit" {
			packetSizeByteArr := make([]byte, 3)

			if _, err := conn.Read(packetSizeByteArr); err != nil {
				conn.Close()
				return err
			}

			packetSize := core.Int24ToInt32(packetSizeByteArr)
			encryptedPacket := make([]byte, packetSize)

			packetSizeInt := int(packetSize)

			totalPositionRead := 0

			for packetSizeInt != totalPositionRead {
				currentPosition, err := conn.Read(encryptedPacket[totalPositionRead:packetSizeInt])
				totalPositionRead += currentPosition

				if err != nil {
					conn.Close()
					return err
				}
			}

			packet, err := bismuth.decryptMessage(aead, encryptedPacket)

			if err != nil {
				conn.Close()
				return err
			}

			if packet[0] == core.InitiateForwarding {
				bmConn := core.BismuthConn{
					Aead:       aead,
					PassedConn: conn,
					MaxBufSize: core.ConnStandardMaxBufSize,
				}

				bmConn.DoInitSteps()

				err := bismuth.HandleConnection(core.BismuthConnWrapped{
					Bismuth: &bmConn,
				})

				return err
			}
		}
	}
}