chatter.go

// Implementation of a forward-secure, end-to-end encrypted messaging client
// supporting key compromise recovery and out-of-order message delivery.
// Directly inspired by Signal/Double-ratchet protocol but missing a few
// features. No asynchronous handshake support (pre-keys) for example.
//
// SECURITY WARNING: This code is meant for educational purposes and may
// contain vulnerabilities or other bugs. Please do not use it for
// security-critical applications.
//
// GRADING NOTES: This is the only file you need to modify for this assignment.
// You may add additional support files if desired. You should modify this file
// to implement the intended protocol, but preserve the function signatures
// for the following methods to ensure your implementation will work with
// standard test code:
//
// *NewChatter
// *EndSession
// *InitiateHandshake
// *ReturnHandshake
// *FinalizeHandshake
// *SendMessage
// *ReceiveMessage
//
// In addition, you'll need to keep all of the following structs' fields:
//
// *Chatter
// *Session
// *Message
//
// You may add fields if needed (not necessary) but don't rename or delete
// any existing fields.
//
// Original version
// Joseph Bonneau February 2019

package chatterbox

import (
	"encoding/binary"
	"errors"
	"fmt"
	"reflect"
	//"bytes"
)

// Labels for key derivation

// Label for generating a check key from the initial root.
// Used for verifying the results of a handshake out-of-band.
const HANDSHAKE_CHECK_LABEL byte = 0x11

// Label for ratcheting the root key after deriving a key chain from it
const ROOT_LABEL = 0x22

// Label for ratcheting the main chain of keys
const CHAIN_LABEL = 0x33

// Label for deriving message keys from chain keys
const KEY_LABEL = 0x44

// Chatter represents a chat participant. Each Chatter has a single long-term
// key Identity, and a map of open sessions with other users (indexed by their
// identity keys). You should not need to modify this.
type Chatter struct {
	Identity *KeyPair
	Sessions map[PublicKey]*Session
}

// Session represents an open session between one chatter and another.
// You should not need to modify this, though you can add additional fields
// if you want to.
type Session struct {
	MyDHRatchet       *KeyPair
	PartnerDHRatchet  *PublicKey
	RootChain         *SymmetricKey
	SendChain         *SymmetricKey
	ReceiveChain      *SymmetricKey
	CachedReceiveKeys map[int]*SymmetricKey
	updateCount       int
	init              int
	LastSender        int
	SendCounter       int
	LastUpdate        int
	ReceiveCounter    int
}

// Message represents a message as sent over an untrusted network.
// The first 5 fields are send unencrypted (but should be authenticated).
// The ciphertext contains the (encrypted) communication payload.
// You should not need to modify this.
type Message struct {
	Sender        *PublicKey
	Receiver      *PublicKey
	NextDHRatchet *PublicKey
	Counter       int
	LastUpdate    int
	Ciphertext    []byte
	IV            []byte
}

// EncodeAdditionalData encodes all of the non-ciphertext fields of a message
// into a single byte array, suitable for use as additional authenticated data
// in an AEAD scheme. You should not need to modify this code.
func (m *Message) EncodeAdditionalData() []byte {
	buf := make([]byte, 8+3*FINGERPRINT_LENGTH)

	binary.LittleEndian.PutUint32(buf, uint32(m.Counter))
	binary.LittleEndian.PutUint32(buf[4:], uint32(m.LastUpdate))

	if m.Sender != nil {
		copy(buf[8:], m.Sender.Fingerprint())
	}
	if m.Receiver != nil {
		copy(buf[8+FINGERPRINT_LENGTH:], m.Receiver.Fingerprint())
	}
	if m.NextDHRatchet != nil {
		copy(buf[8+2*FINGERPRINT_LENGTH:], m.NextDHRatchet.Fingerprint())
	}

	return buf
}

// NewChatter creates and initializes a new Chatter object. A long-term
// identity key is created and the map of sessions is initialized.
// You should not need to modify this code.
func NewChatter() *Chatter {
	c := new(Chatter)
	c.Identity = GenerateKeyPair()
	c.Sessions = make(map[PublicKey]*Session)
	return c
}

// EndSession erases all data for a session with the designated partner.
// All outstanding key material should be zeroized and the session erased.
func (c *Chatter) EndSession(partnerIdentity *PublicKey) error {

	if _, exists := c.Sessions[*partnerIdentity]; !exists {
		return errors.New("Don't have that session open to tear down")
	}

	/*MyDHRatchet       *KeyPair
	PartnerDHRatchet  *PublicKey
	RootChain         *SymmetricKey
	SendChain         *SymmetricKey
	ReceiveChain      *SymmetricKey
	CachedReceiveKeys map[int]*SymmetricKey
	init              int
	LastSender        int
	SendCounter       int
	LastUpdate        int
	ReceiveCounter    int*/

	// TODO: your code here to zeroize remaining state
	c.Sessions[*partnerIdentity].MyDHRatchet.Zeroize()
	c.Sessions[*partnerIdentity].PartnerDHRatchet = nil
	c.Sessions[*partnerIdentity].RootChain.Zeroize()
	c.Sessions[*partnerIdentity].SendChain.Zeroize()
	c.Sessions[*partnerIdentity].ReceiveChain.Zeroize()
	c.Sessions[*partnerIdentity].init = 0
	c.Sessions[*partnerIdentity].LastSender = 0
	c.Sessions[*partnerIdentity].SendCounter = 0
	c.Sessions[*partnerIdentity].LastSender = 0
	c.Sessions[*partnerIdentity].ReceiveCounter = 0
	for k, v := range c.Sessions[*partnerIdentity].CachedReceiveKeys {
		defer delete(c.Sessions[*partnerIdentity].CachedReceiveKeys, k)
		v.Zeroize()
	}
	delete(c.Sessions, *partnerIdentity)

	return nil
}

// InitiateHandshake prepares the first message sent in a handshake, containing
// an ephemeral DH share. The partner which calls this method is the initiator.
func (c *Chatter) InitiateHandshake(partnerIdentity *PublicKey) (*PublicKey, error) {

	if _, exists := c.Sessions[*partnerIdentity]; exists {
		return nil, errors.New("Already have session open")
	}

	c.Sessions[*partnerIdentity] = &Session{
		CachedReceiveKeys: make(map[int]*SymmetricKey),
		// TODO: your code here
		MyDHRatchet:      nil,
		PartnerDHRatchet: nil,
		RootChain:        nil,
		SendChain:        nil,
		ReceiveChain:     nil,
		updateCount:      0,
		init:             1,
		LastSender:       0,
		SendCounter:      0,
		LastUpdate:       0,
		ReceiveCounter:   0,
	}

	// TODO: your code here
	var ephemeralKeys = GenerateKeyPair()
	c.Sessions[*partnerIdentity].MyDHRatchet = ephemeralKeys.Duplicate()

	return &ephemeralKeys.PublicKey, nil
}

// ReturnHandshake prepares the second message sent in a handshake, containing
// an ephemeral DH share. The partner which calls this method is the responder.
func (c *Chatter) ReturnHandshake(partnerIdentity,
	partnerEphemeral *PublicKey) (*PublicKey, *SymmetricKey, error) {

	if _, exists := c.Sessions[*partnerIdentity]; exists {
		return nil, nil, errors.New("Already have session open")
	}

	c.Sessions[*partnerIdentity] = &Session{
		CachedReceiveKeys: make(map[int]*SymmetricKey),
		// TODO: your code here
		MyDHRatchet:      nil,
		PartnerDHRatchet: nil,
		RootChain:        nil,
		SendChain:        nil,
		ReceiveChain:     nil,
		updateCount:      0,
		init:             0,
		SendCounter:      0,
		LastSender:       0,
		LastUpdate:       0,
		ReceiveCounter:   0,
	}

	// TODO: your code here
	var selfEphemeral = GenerateKeyPair()
	//g A(Alice public key)b(self ephemeral key)
	var gAb = DHCombine(partnerIdentity, &selfEphemeral.PrivateKey)
	//g a(Alice ephemeral key) B(self long-term private key)
	var gBa = DHCombine(partnerEphemeral, &c.Identity.PrivateKey)
	//g a(Alice ephemeral) b(self ephemeral)
	var gab = DHCombine(partnerEphemeral, &selfEphemeral.PrivateKey)

	var combinedKey = CombineKeys(gAb, gBa, gab)
	var rootkey = combinedKey.DeriveKey(HANDSHAKE_CHECK_LABEL)
	c.Sessions[*partnerIdentity].MyDHRatchet = selfEphemeral.Duplicate()
	c.Sessions[*partnerIdentity].PartnerDHRatchet = partnerEphemeral.Duplicate()
	c.Sessions[*partnerIdentity].RootChain = combinedKey.Duplicate()
	c.Sessions[*partnerIdentity].ReceiveChain = c.Sessions[*partnerIdentity].RootChain.DeriveKey(CHAIN_LABEL).Duplicate()
	return &selfEphemeral.PublicKey, rootkey, nil
}

// FinalizeHandshake lets the initiator receive the responder's ephemeral key
// and finalize the handshake.The partner which calls this method is the initiator.
func (c *Chatter) FinalizeHandshake(partnerIdentity,
	partnerEphemeral *PublicKey) (*SymmetricKey, error) {

	if _, exists := c.Sessions[*partnerIdentity]; !exists {
		return nil, errors.New("Can't finalize session, not yet open")
	}

	// TODO: your code here
	//gAb
	var gAb = DHCombine(partnerEphemeral, &c.Identity.PrivateKey)
	//gBa
	var gBa = DHCombine(partnerIdentity, &c.Sessions[*partnerIdentity].MyDHRatchet.PrivateKey)
	//gab
	var gab = DHCombine(partnerEphemeral, &c.Sessions[*partnerIdentity].MyDHRatchet.PrivateKey)

	var combinedKey = CombineKeys(gAb, gBa, gab)
	var rootkey = combinedKey.DeriveKey(HANDSHAKE_CHECK_LABEL)
	c.Sessions[*partnerIdentity].PartnerDHRatchet = partnerEphemeral.Duplicate()
	c.Sessions[*partnerIdentity].RootChain = combinedKey.Duplicate()

	return rootkey, nil
}

// SendMessage is used to send the given plaintext string as a message.
// You'll need to implement the code to ratchet, derive keys and encrypt this message.
func (c *Chatter) SendMessage(partnerIdentity *PublicKey,
	plaintext string) (*Message, error) {

	if _, exists := c.Sessions[*partnerIdentity]; !exists {
		return nil, errors.New("Can't send message to partner with no open session")
	}

	message := &Message{
		Sender:        &c.Identity.PublicKey,
		Receiver:      partnerIdentity,
		NextDHRatchet: nil,
		Counter:       0,
		LastUpdate:    0,
		Ciphertext:    nil,
		IV:            nil,
	}

	// TODO: your code here

	switch {
	//Alice First Time
	case (c.Sessions[*partnerIdentity].SendChain == nil) && (c.Sessions[*partnerIdentity].init == 1) && (c.Sessions[*partnerIdentity].SendCounter == 0):

		c.Sessions[*partnerIdentity].SendChain = c.Sessions[*partnerIdentity].RootChain.DeriveKey(CHAIN_LABEL)

		c.Sessions[*partnerIdentity].SendCounter++

		c.Sessions[*partnerIdentity].LastUpdate = 0

		c.Sessions[*partnerIdentity].LastSender = 1

	//Bob first time
	case (c.Sessions[*partnerIdentity].SendChain == nil) && (c.Sessions[*partnerIdentity].init == 0):
		fmt.Print("\nBOB FIRST SEND")
		var newEphDH = GenerateKeyPair()

		var gab = DHCombine(c.Sessions[*partnerIdentity].PartnerDHRatchet, &newEphDH.PrivateKey)

		//var cleanLocation = &c.Sessions[*partnerIdentity].MyDHRatchet

		c.Sessions[*partnerIdentity].MyDHRatchet = newEphDH.Duplicate()

		//(*cleanLocation).Zeroize()

		c.Sessions[*partnerIdentity].RootChain = c.Sessions[*partnerIdentity].RootChain.Duplicate().DeriveKey(ROOT_LABEL)

		c.Sessions[*partnerIdentity].RootChain = CombineKeys(c.Sessions[*partnerIdentity].RootChain, gab).Duplicate()

		fmt.Print("\nSENT ROOTCHAIN BOB: ", c.Sessions[*partnerIdentity].RootChain, "\n")
		c.Sessions[*partnerIdentity].SendChain = c.Sessions[*partnerIdentity].RootChain.DeriveKey(CHAIN_LABEL)

		c.Sessions[*partnerIdentity].SendCounter++

		c.Sessions[*partnerIdentity].LastUpdate = c.Sessions[*partnerIdentity].SendCounter

		message.LastUpdate = c.Sessions[*partnerIdentity].SendCounter

		c.Sessions[*partnerIdentity].updateCount++

		c.Sessions[*partnerIdentity].LastSender = 1

		message.NextDHRatchet = &newEphDH.PublicKey

	//If the sender is not the sender of the last msg, then perform DH rachet
	case c.Sessions[*partnerIdentity].LastSender == 0:

		fmt.Print("\n DIFFERNT SENDER \n")

		var newEphDH = GenerateKeyPair()

		var gab = DHCombine(c.Sessions[*partnerIdentity].PartnerDHRatchet, &newEphDH.PrivateKey)

		fmt.Print("\n PARTNERDHRATCHET(SHOULD BE SAME AS LAST): ", c.Sessions[*partnerIdentity].PartnerDHRatchet)

		//var cleanLocation = c.Sessions[*partnerIdentity].MyDHRatchet

		c.Sessions[*partnerIdentity].MyDHRatchet = newEphDH.Duplicate()

		//(*cleanLocation).Zeroize()

		//RATCHET ROOT

		//var cleanRoot = &c.Sessions[*partnerIdentity].RootChain

		c.Sessions[*partnerIdentity].RootChain = c.Sessions[*partnerIdentity].RootChain.Duplicate().DeriveKey(ROOT_LABEL)

		//(*cleanRoot).Zeroize()

		c.Sessions[*partnerIdentity].RootChain = CombineKeys(c.Sessions[*partnerIdentity].RootChain, gab).Duplicate()

		c.Sessions[*partnerIdentity].SendChain = c.Sessions[*partnerIdentity].RootChain.DeriveKey(CHAIN_LABEL)

		fmt.Print("\nSENT ROOTCHAIN: ", c.Sessions[*partnerIdentity].RootChain, "\n")

		c.Sessions[*partnerIdentity].SendCounter++

		c.Sessions[*partnerIdentity].LastUpdate = c.Sessions[*partnerIdentity].SendCounter

		message.LastUpdate = c.Sessions[*partnerIdentity].SendCounter

		c.Sessions[*partnerIdentity].updateCount++

		c.Sessions[*partnerIdentity].LastSender = 1

		message.NextDHRatchet = &newEphDH.PublicKey
	//else just Rachet normally
	case c.Sessions[*partnerIdentity].LastSender == 1:
		var nextSend = c.Sessions[*partnerIdentity].SendChain.DeriveKey(CHAIN_LABEL)

		//var cleanLocation = &c.Sessions[*partnerIdentity].SendChain

		c.Sessions[*partnerIdentity].SendChain = nextSend.Duplicate()

		//(*cleanLocation).Zeroize()

		//nextSend.Zeroize()

		c.Sessions[*partnerIdentity].SendCounter++
		message.LastUpdate = c.Sessions[*partnerIdentity].LastUpdate
		message.NextDHRatchet = &c.Sessions[*partnerIdentity].MyDHRatchet.PublicKey
	}
	//var cleanLocation = &c.Sessions[*partnerIdentity].SendChain
	//c.Sessions[*partnerIdentity].SendChain = c.Sessions[*partnerIdentity].SendChain.DeriveKey(CHAIN_LABEL)
	//(*cleanLocation).Zeroize()
	message.Counter = c.Sessions[*partnerIdentity].SendCounter
	fmt.Print("\nSENT DHRACHET: ", message.NextDHRatchet, "\n")
	fmt.Print("\nSENT SENDCHAIN: ", c.Sessions[*partnerIdentity].SendChain, "\n")
	var msgKey = c.Sessions[*partnerIdentity].SendChain.DeriveKey(KEY_LABEL)
	fmt.Print("\nSENT MSG KEY: ", msgKey, "\n MSG COUNT: ", message.Counter, "\n SENT BY: ", c.Sessions[*partnerIdentity].init, "\n LAST UPDATE: ", message.LastUpdate, "\n")
	var niv = NewIV()
	message.IV = niv
	var additionalData = message.EncodeAdditionalData()
	message.Ciphertext = msgKey.AuthenticatedEncrypt(plaintext, additionalData, niv)
	return message, nil
}

// ReceiveMessage is used to receive the given message and return the correct
// plaintext. This method is where most of the key derivation, ratcheting
// and out-of-order message handling logic happens.
func (c *Chatter) ReceiveMessage(message *Message) (string, error) {

	if _, exists := c.Sessions[*message.Sender]; !exists {
		return "", errors.New("Can't receive message from partner with no open session")
	}

	// TODO: your code here
	var flag = 0
	if message.Counter != c.Sessions[*message.Sender].ReceiveCounter+1 {
		if c.Sessions[*message.Sender].ReceiveChain == nil {
			fmt.Println("\n RECEIVECHAIN WAS NIL")
			c.Sessions[*message.Sender].ReceiveChain = c.Sessions[*message.Sender].RootChain.DeriveKey(CHAIN_LABEL)

		}
		//fmt.Print("\n WE ARE LOOPING \n")
		c.DeriveKeysTillRecent(message)
		flag = 1
	} else {

		if (!reflect.DeepEqual(c.Sessions[*message.Sender].PartnerDHRatchet, message.NextDHRatchet)) && (message.NextDHRatchet != nil) {
			c.NewkeyRatchet(message)
		} else if message.NextDHRatchet != nil {
			c.RegularRatchet(message)
		}
		flag = 1
	}
	fmt.Println("\n RECEIVECHAIN: ", c.Sessions[*message.Sender].ReceiveChain)
	/*if c.Sessions[*message.Sender].ReceiveChain != nil {
		flag = 0
	}*/

	var additionalData = message.EncodeAdditionalData()
	if c.Sessions[*message.Sender].ReceiveChain == nil {
		fmt.Println("\n RECEIVECHAIN WAS NIL")
		c.Sessions[*message.Sender].ReceiveChain = c.Sessions[*message.Sender].RootChain.DeriveKey(CHAIN_LABEL)
	} else if flag == 0 {
		//REGULAR RACHET
		c.RegularRatchet(message)
	}

	var msgKey = c.Sessions[*message.Sender].ReceiveChain.DeriveKey(KEY_LABEL).Duplicate()

	if (!reflect.DeepEqual(c.Sessions[*message.Sender].PartnerDHRatchet, message.NextDHRatchet)) && (c.Sessions[*message.Sender].ReceiveCounter > message.Counter) {
		msgKey = c.Sessions[*message.Sender].CachedReceiveKeys[message.Counter]
	} else {
		msgKey = c.Sessions[*message.Sender].ReceiveChain.DeriveKey(KEY_LABEL)
		c.Sessions[*message.Sender].ReceiveCounter = message.Counter
	}

	//set LastSender = 0 for Alice if msg lastUpdate > self LastUpdate, set LastSender = 0 for Bob if msg lastUpdate => self LastUpdate

	fmt.Print("\nRECIEVED msgKey: ", msgKey, "\n MSG COUNT: ", message.Counter, "\n RECEIVED BY: ", c.Sessions[*message.Sender].init, "\n LAST UPDATE: ", message.LastUpdate, "NEW RECEIVE COUNTER:", c.Sessions[*message.Sender].ReceiveCounter, "\n")
	//msgKey.Zeroize()
	var msgOut, err = msgKey.AuthenticatedDecrypt(message.Ciphertext, additionalData, message.IV)
	fmt.Print("ERROR: ", err)
	fmt.Print("\n===================================================================================================================\n")
	return msgOut, err
}

// HELPER FUNCTIONS
func (c *Chatter) NewkeyRatchet(message *Message) {
	fmt.Println("DEBUG: ", message.NextDHRatchet)
	var gab = DHCombine(message.NextDHRatchet, &c.Sessions[*message.Sender].MyDHRatchet.PrivateKey)
	fmt.Print("\n RECIVE DHRATCHET PRIVATE KEY: ", c.Sessions[*message.Sender].RootChain, "\n")
	var ratchetRoot = c.Sessions[*message.Sender].RootChain.Duplicate().DeriveKey(ROOT_LABEL)
	c.Sessions[*message.Sender].RootChain = CombineKeys(ratchetRoot, gab)
	fmt.Print("\n RECIVE ROOT CHAIN: ", c.Sessions[*message.Sender].RootChain, "\n")
	c.Sessions[*message.Sender].PartnerDHRatchet = message.NextDHRatchet.Duplicate()
	fmt.Print("\n RECIVE PARTNER DHR: ", c.Sessions[*message.Sender].PartnerDHRatchet, "\n")
	c.Sessions[*message.Sender].ReceiveChain = nil
	c.Sessions[*message.Sender].LastSender = 0
}
func (c *Chatter) RegularRatchet(message *Message) {
	//var cleanLocation = &c.Sessions[*message.Sender].ReceiveChain
	fmt.Print("\n REGULAR RACHET, RECEIVE: ", c.Sessions[*message.Sender].ReceiveChain, "\n")
	c.Sessions[*message.Sender].ReceiveChain = c.Sessions[*message.Sender].ReceiveChain.DeriveKey(CHAIN_LABEL)
	//defer (*cleanLocation).Zeroize()
	c.Sessions[*message.Sender].LastSender = 0
}

func (c *Chatter) DeriveKeysTillRecent(message *Message) {
	var currentReciveCount = c.Sessions[*message.Sender].ReceiveCounter
	var msgnumber = message.Counter
	var rachetTime = message.LastUpdate
	fmt.Print("RACHETTIME: ", rachetTime)
	if currentReciveCount < msgnumber {
		for i := currentReciveCount + 1; i < msgnumber; i++ {

			if (i == rachetTime) && (message.NextDHRatchet != nil) {
				c.NewkeyRatchet(message)
				if c.Sessions[*message.Sender].ReceiveChain == nil {
					fmt.Println("\n RECEIVECHAIN WAS NIL IN LOOP")
					c.Sessions[*message.Sender].ReceiveChain = c.Sessions[*message.Sender].RootChain.DeriveKey(CHAIN_LABEL)
				}
			}
			var newMsgKey = c.Sessions[*message.Sender].ReceiveChain.DeriveKey(KEY_LABEL)
			c.Sessions[*message.Sender].CachedReceiveKeys[i] = newMsgKey.Duplicate()
			newMsgKey.Zeroize()
			fmt.Print("CACHED MSG KEY: ", i, " : ", c.Sessions[*message.Sender].CachedReceiveKeys[i])
			c.RegularRatchet(message)
			fmt.Print("\n REGULAR RACHET, NEW KEY: ", c.Sessions[*message.Sender].ReceiveChain, "\n\n\n")
		}
	}

}