From a06e714c3ce9c1a25c26dea5d6e334f368e9afc7 Mon Sep 17 00:00:00 2001
From: Oleg Kovalov <oleg@hey.com>
Date: Thu, 10 Aug 2023 20:19:23 +0200
Subject: [PATCH] Add ObjectID (#1)

---
 objectid.go      | 87 ++++++++++++++++++++++++++++++++++++++++++++++++
 objectid_test.go | 11 ++++++
 utils.go         |  8 +++++
 utils_test.go    | 11 ++++++
 4 files changed, 117 insertions(+)
 create mode 100644 objectid.go
 create mode 100644 objectid_test.go
 create mode 100644 utils.go
 create mode 100644 utils_test.go

diff --git a/objectid.go b/objectid.go
new file mode 100644
index 0000000..5264d0e
--- /dev/null
+++ b/objectid.go
@@ -0,0 +1,87 @@
+package bson
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	mathrand "math/rand"
+	"sync/atomic"
+	"time"
+)
+
+var (
+	_ Marshaler   = &ObjectID{}
+	_ Unmarshaler = &ObjectID{}
+)
+
+var ErrBadObjectID = errors.New("provided input is not a valid ObjectID")
+
+// ObjectID represents BSON object ID.
+type ObjectID [12]byte
+
+// NewObjectID returns a new ObjectID.
+func NewObjectID() ObjectID {
+	return NewObjectIDWithTime(time.Now())
+}
+
+// NewObjectIDWithTime returns a new ObjectID.
+func NewObjectIDWithTime(t time.Time) ObjectID {
+	ts := uint32(t.UTC().Unix())
+	c := objectIDCounter.Add(1)
+
+	var oid ObjectID
+	oid[0] = byte(ts >> 24)
+	oid[1] = byte(ts >> 16)
+	oid[2] = byte(ts >> 8)
+	oid[3] = byte(ts)
+
+	oid[4] = procUniqueID[0]
+	oid[5] = procUniqueID[1]
+	oid[6] = procUniqueID[2]
+	oid[7] = procUniqueID[3]
+	oid[8] = procUniqueID[4]
+
+	oid[9] = byte(c >> 16)
+	oid[10] = byte(c >> 8)
+	oid[11] = byte(c)
+	return oid
+}
+
+// String returns a hex string representation of the id.
+// Example: ObjectIdHex('64d526fa37931c1e97eea90f').
+func (oid ObjectID) String() string {
+	return fmt.Sprintf(`ObjectIdHex('%x')`, string(oid[:]))
+}
+
+func (oid *ObjectID) MarshalBSON() ([]byte, error) {
+	b := make([]byte, len(oid))
+	copy(b, oid[:])
+	return b, nil
+}
+
+func (oid *ObjectID) UnmarshalBSON(b []byte) error {
+	switch len(b) {
+	case 12:
+		copy(oid[:], b)
+		return nil
+	case 24:
+		n, err := hex.Decode(oid[:], b)
+		if n != 24 {
+			panic("unreachable")
+		}
+		return err
+	default:
+		return ErrBadObjectID
+	}
+}
+
+var (
+	procUniqueID    [5]byte
+	objectIDCounter atomic.Uint32
+)
+
+func init() {
+	must(rand.Read(procUniqueID[:]))
+	objectIDCounter.Store(mathrand.Uint32())
+}
diff --git a/objectid_test.go b/objectid_test.go
new file mode 100644
index 0000000..1416c90
--- /dev/null
+++ b/objectid_test.go
@@ -0,0 +1,11 @@
+package bson
+
+import (
+	"testing"
+)
+
+func TestObjectID(t *testing.T) {
+	oid := ObjectID([12]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})
+
+	mustEqual(t, oid.String(), "ObjectIdHex('0102030405060708090a0b0c')")
+}
diff --git a/utils.go b/utils.go
new file mode 100644
index 0000000..6cb8ef8
--- /dev/null
+++ b/utils.go
@@ -0,0 +1,8 @@
+package bson
+
+func must[T any](v T, err error) T {
+	if err != nil {
+		panic(err)
+	}
+	return v
+}
diff --git a/utils_test.go b/utils_test.go
new file mode 100644
index 0000000..011b0bd
--- /dev/null
+++ b/utils_test.go
@@ -0,0 +1,11 @@
+package bson
+
+import "testing"
+
+func mustEqual[T comparable](tb testing.TB, have, want T) {
+	tb.Helper()
+
+	if have != want {
+		tb.Fatalf("\nhave: %+v\nwant: %+v\n", have, want)
+	}
+}