From 4ccf95bee03bc558bb4feaba6443da10d0d2a42c Mon Sep 17 00:00:00 2001
From: Sarah Jamie Lewis <me@sarahjamielewis.com>
Date: Fri, 12 Jan 2018 13:02:15 -0500
Subject: [PATCH] Adding RandNumber util to simplify chatchannel logic

---
 channels/chatchannel.go      | 13 ++-----------
 channels/chatchannel_test.go |  2 ++
 utils/crypto.go              | 10 ++++++++++
 utils/crypto_test.go         |  8 ++++++++
 4 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/channels/chatchannel.go b/channels/chatchannel.go
index 01ba01e..001ad46 100644
--- a/channels/chatchannel.go
+++ b/channels/chatchannel.go
@@ -1,13 +1,10 @@
 package channels
 
 import (
-	"crypto/rand"
 	"github.com/golang/protobuf/proto"
 	"github.com/s-rah/go-ricochet/utils"
 	"github.com/s-rah/go-ricochet/wire/chat"
 	"github.com/s-rah/go-ricochet/wire/control"
-	"math"
-	"math/big"
 	"time"
 )
 
@@ -100,10 +97,7 @@ func (cc *ChatChannel) RequiresAuthentication() string {
 // returned, it will be sent as the ChannelResult message.
 func (cc *ChatChannel) OpenInbound(channel *Channel, raw *Protocol_Data_Control.OpenChannel) ([]byte, error) {
 	cc.channel = channel
-	id, err := rand.Int(rand.Reader, big.NewInt(math.MaxUint32))
-	if err != nil {
-		return nil, err
-	}
+	id := utils.GetRandNumber()
 	cc.lastMessageID = uint32(id.Uint64())
 	cc.channel.Pending = false
 	messageBuilder := new(utils.MessageBuilder)
@@ -115,10 +109,7 @@ func (cc *ChatChannel) OpenInbound(channel *Channel, raw *Protocol_Data_Control.
 // returned, it will be sent as the OpenChannel message.
 func (cc *ChatChannel) OpenOutbound(channel *Channel) ([]byte, error) {
 	cc.channel = channel
-	id, err := rand.Int(rand.Reader, big.NewInt(math.MaxUint32))
-	if err != nil {
-		return nil, err
-	}
+	id := utils.GetRandNumber()
 	cc.lastMessageID = uint32(id.Uint64())
 	messageBuilder := new(utils.MessageBuilder)
 	return messageBuilder.OpenChannel(channel.ID, cc.Type()), nil
diff --git a/channels/chatchannel_test.go b/channels/chatchannel_test.go
index 8ea0b12..c34c5a2 100644
--- a/channels/chatchannel_test.go
+++ b/channels/chatchannel_test.go
@@ -120,5 +120,7 @@ func TestChatChannelOperations(t *testing.T) {
 	chatChannel.Packet(chat)
 
 	chatChannel.SendMessage("hello")
+	ackMessage := messageBuilder.AckChatMessage(0, true)
+        chatChannel.Packet(ackMessage)
 
 }
diff --git a/utils/crypto.go b/utils/crypto.go
index cdf3537..f8366f0 100644
--- a/utils/crypto.go
+++ b/utils/crypto.go
@@ -7,6 +7,8 @@ import (
 	"encoding/pem"
 	"errors"
 	"io/ioutil"
+	"math"
+	"math/big"
 )
 
 const (
@@ -17,6 +19,14 @@ const (
 	RicochetKeySize = 1024
 )
 
+func GetRandNumber() *big.Int {
+        num, err := rand.Int(rand.Reader, big.NewInt(math.MaxUint32))
+        // If we can't generate random numbers then panicking is probably
+        // the best option.
+        CheckError(err)
+        return num
+}
+
 // GeneratePrivateKey generates a new private key for use
 func GeneratePrivateKey() (*rsa.PrivateKey, error) {
 	privateKey, err := rsa.GenerateKey(rand.Reader, RicochetKeySize)
diff --git a/utils/crypto_test.go b/utils/crypto_test.go
index 2436dae..88ef735 100644
--- a/utils/crypto_test.go
+++ b/utils/crypto_test.go
@@ -2,6 +2,7 @@ package utils
 
 import (
 	"testing"
+	"math"
 )
 
 func TestGeneratePrivateKey(t *testing.T) {
@@ -17,3 +18,10 @@ func TestLoadPrivateKey(t *testing.T) {
 		t.Errorf("Error while loading private key from file: %v", err)
 	}
 }
+
+func TestGetRandNumber(t *testing.T) {
+        num := GetRandNumber()
+        if !num.IsUint64() || num.Uint64() > uint64(math.MaxUint32) {
+                t.Errorf("Error random number outside of expected bounds %v", num)
+        }
+}