add optional order flag

cmd/discreteLog.go

@@ -23,10 +23,10 @@ import (
 	"github.com/spf13/cobra"
 )
 
-// FIXME: order and generator may not be the right names here
 var discreteLogModulus string
 var discreteLogGenerator string
 var discreteLogElement string
+var discreteLogOrder string
 
 // whyyyy doesn't math/big have a ceil functionnnn
 func ceilSqrt(x *big.Int) *big.Int {
@@ -39,9 +39,16 @@ func ceilSqrt(x *big.Int) *big.Int {
 	return z
 }
 
-// FIXME: n is supposed to be the order, but i'm using it as the modulus
-func babyStepGiantStep(n, g, x *big.Int) *big.Int {
-	m := ceilSqrt(n)
+func babyStepGiantStep(n, g, x, order *big.Int) *big.Int {
+	var m *big.Int
+	if order == nil {
+		// m = ceil(sqrt(n - 1))
+		z := big.NewInt(1)
+		z.Sub(n, z)
+		m = ceilSqrt(z)
+	} else {
+		m = ceilSqrt(order)
+	}
 
 	table := make(map[string]*big.Int)
 	for j := big.NewInt(1); j.Cmp(m) <= 0; j.Add(j, big.NewInt(1)) {
@@ -87,7 +94,15 @@ func discreteLog(cmd *cobra.Command, args []string) {
 		cobra.CheckErr("invalid element " + discreteLogElement)
 	}
 
-	k := babyStepGiantStep(n, g, x)
+	var order *big.Int
+	if discreteLogOrder != "" {
+		order, ok = new(big.Int).SetString(discreteLogOrder, 10)
+		if !ok {
+			cobra.CheckErr("invalid order " + discreteLogOrder)
+		}
+	}
+
+	k := babyStepGiantStep(n, g, x, order)
 	fmt.Println(k)
 }
 
@@ -95,7 +110,7 @@ func discreteLog(cmd *cobra.Command, args []string) {
 var discreteLogCmd = &cobra.Command{
 	Use:   "discrete-log",
 	Short: "Compute the discrete logarithm",
-	Long:  `Compute the discrete logarithm`,
+	Long:  `Compute the discrete logarithm.`,
 	Run:   discreteLog,
 }
 
@@ -120,4 +135,6 @@ func init() {
 
 	discreteLogCmd.Flags().StringVarP(&discreteLogElement, "element", "e", "", "element of the cyclic group to compute logarithm of")
 	discreteLogCmd.MarkFlagRequired("element")
+
+	discreteLogCmd.Flags().StringVar(&discreteLogOrder, "order", "", "order of the cyclic group")
 }