2 files changed, 183 insertions, 3 deletions

diff --git a/sshbuf-getput-basic.c b/sshbuf-getput-basic.c
index 50648258f..27058d5bb 100644
--- a/sshbuf-getput-basic.c
+++ b/sshbuf-getput-basic.c
@@ -1,4 +1,4 @@
-/*      $OpenBSD: sshbuf-getput-basic.c,v 1.7 2017/06/01 04:51:58 djm Exp $     */
+/*      $OpenBSD: sshbuf-getput-basic.c,v 1.8 2019/07/14 23:32:27 djm Exp $     */
 /*
  * Copyright (c) 2011 Damien Miller
  *
@@ -93,6 +93,93 @@ sshbuf_get_u8(struct sshbuf *buf, u_char *valp)
         return 0;
 }
 
+static int
+check_offset(const struct sshbuf *buf, int wr, size_t offset, size_t len)
+{
+        if (sshbuf_ptr(buf) == NULL) /* calls sshbuf_check_sanity() */
+                return SSH_ERR_INTERNAL_ERROR;
+        if (offset >= SIZE_MAX - len)
+                return SSH_ERR_INVALID_ARGUMENT;
+        if (offset + len > sshbuf_len(buf)) {
+                return wr ?
+                    SSH_ERR_NO_BUFFER_SPACE : SSH_ERR_MESSAGE_INCOMPLETE;
+        }
+        return 0;
+}
+
+static int
+check_roffset(const struct sshbuf *buf, size_t offset, size_t len,
+    const u_char **p)
+{
+        int r;
+
+        *p = NULL;
+        if ((r = check_offset(buf, 0, offset, len)) != 0)
+                return r;
+        *p = sshbuf_ptr(buf) + offset;
+        return 0;
+}
+
+int
+sshbuf_peek_u64(const struct sshbuf *buf, size_t offset, u_int64_t *valp)
+{
+        const u_char *p = NULL;
+        int r;
+
+        if (valp != NULL)
+                *valp = 0;
+        if ((r = check_roffset(buf, offset, 8, &p)) != 0)
+                return r;
+        if (valp != NULL)
+                *valp = PEEK_U64(p);
+        return 0;
+}
+
+int
+sshbuf_peek_u32(const struct sshbuf *buf, size_t offset, u_int32_t *valp)
+{
+        const u_char *p = NULL;
+        int r;
+
+        if (valp != NULL)
+                *valp = 0;
+        if ((r = check_roffset(buf, offset, 4, &p)) != 0)
+                return r;
+        if (valp != NULL)
+                *valp = PEEK_U32(p);
+        return 0;
+}
+
+int
+sshbuf_peek_u16(const struct sshbuf *buf, size_t offset, u_int16_t *valp)
+{
+        const u_char *p = NULL;
+        int r;
+
+        if (valp != NULL)
+                *valp = 0;
+        if ((r = check_roffset(buf, offset, 2, &p)) != 0)
+                return r;
+        if (valp != NULL)
+                *valp = PEEK_U16(p);
+        return 0;
+}
+
+int
+sshbuf_peek_u8(const struct sshbuf *buf, size_t offset, u_char *valp)
+{
+        const u_char *p = NULL;
+        int r;
+
+        if (valp != NULL)
+                *valp = 0;
+        if ((r = check_roffset(buf, offset, 1, &p)) != 0)
+                return r;
+        if (valp != NULL)
+                *valp = *p;
+        return 0;
+}
+
 int
 sshbuf_get_string(struct sshbuf *buf, u_char **valp, size_t *lenp)
 {
@@ -344,6 +431,80 @@ sshbuf_put_u8(struct sshbuf *buf, u_char val)
         return 0;
 }
 
+static int
+check_woffset(struct sshbuf *buf, size_t offset, size_t len, u_char **p)
+{
+        int r;
+
+        *p = NULL;
+        if ((r = check_offset(buf, 1, offset, len)) != 0)
+                return r;
+        if (sshbuf_mutable_ptr(buf) == NULL)
+                return SSH_ERR_BUFFER_READ_ONLY;
+        *p = sshbuf_mutable_ptr(buf) + offset;
+        return 0;
+}
+
+int
+sshbuf_poke_u64(struct sshbuf *buf, size_t offset, u_int64_t val)
+{
+        u_char *p = NULL;
+        int r;
+
+        if ((r = check_woffset(buf, offset, 8, &p)) != 0)
+                return r;
+        POKE_U64(p, val);
+        return 0;
+}
+
+int
+sshbuf_poke_u32(struct sshbuf *buf, size_t offset, u_int32_t val)
+{
+        u_char *p = NULL;
+        int r;
+
+        if ((r = check_woffset(buf, offset, 4, &p)) != 0)
+                return r;
+        POKE_U32(p, val);
+        return 0;
+}
+
+int
+sshbuf_poke_u16(struct sshbuf *buf, size_t offset, u_int16_t val)
+{
+        u_char *p = NULL;
+        int r;
+
+        if ((r = check_woffset(buf, offset, 2, &p)) != 0)
+                return r;
+        POKE_U16(p, val);
+        return 0;
+}
+
+int
+sshbuf_poke_u8(struct sshbuf *buf, size_t offset, u_char val)
+{
+        u_char *p = NULL;
+        int r;
+
+        if ((r = check_woffset(buf, offset, 1, &p)) != 0)
+                return r;
+        *p = val;
+        return 0;
+}
+
+int
+sshbuf_poke(struct sshbuf *buf, size_t offset, void *v, size_t len)
+{
+        u_char *p = NULL;
+        int r;
+
+        if ((r = check_woffset(buf, offset, len, &p)) != 0)
+                return r;
+        memcpy(p, v, len);
+        return 0;
+}
+
 int
 sshbuf_put_string(struct sshbuf *buf, const void *v, size_t len)
 {
diff --git a/sshbuf.h b/sshbuf.h
index 7900b82ba..1d9d5bb23 100644
--- a/sshbuf.h
+++ b/sshbuf.h
@@ -1,4 +1,4 @@
-/*      $OpenBSD: sshbuf.h,v 1.13 2019/01/21 09:54:11 djm Exp $ */
+/*      $OpenBSD: sshbuf.h,v 1.14 2019/07/14 23:32:27 djm Exp $ */
 /*
  * Copyright (c) 2011 Damien Miller
  *
@@ -176,6 +176,26 @@ int	sshbuf_put_u32(struct sshbuf *buf, u_int32_t val);
 int     sshbuf_put_u16(struct sshbuf *buf, u_int16_t val);
 int     sshbuf_put_u8(struct sshbuf *buf, u_char val);
 
+/* Functions to peek at the contents of a buffer without modifying it. */
+int     sshbuf_peek_u64(const struct sshbuf *buf, size_t offset,
+    u_int64_t *valp);
+int     sshbuf_peek_u32(const struct sshbuf *buf, size_t offset,
+    u_int32_t *valp);
+int     sshbuf_peek_u16(const struct sshbuf *buf, size_t offset,
+    u_int16_t *valp);
+int     sshbuf_peek_u8(const struct sshbuf *buf, size_t offset,
+    u_char *valp);
+
+/*
+ * Functions to poke values into an exisiting buffer (e.g. a length header
+ * to a packet). The destination bytes must already exist in the buffer.
+ */
+int sshbuf_poke_u64(struct sshbuf *buf, size_t offset, u_int64_t val);
+int sshbuf_poke_u32(struct sshbuf *buf, size_t offset, u_int32_t val);
+int sshbuf_poke_u16(struct sshbuf *buf, size_t offset, u_int16_t val);
+int sshbuf_poke_u8(struct sshbuf *buf, size_t offset, u_char val);
+int sshbuf_poke(struct sshbuf *buf, size_t offset, void *v, size_t len);
+
 /*
  * Functions to extract or store SSH wire encoded strings (u32 len || data)
  * The "cstring" variants admit no \0 characters in the string contents.
@@ -202,7 +222,6 @@ int	sshbuf_get_string_direct(struct sshbuf *buf, const u_char **valp,
 /* Another variant: "peeks" into the buffer without modifying it */
 int     sshbuf_peek_string_direct(const struct sshbuf *buf, const u_char **valp,
             size_t *lenp);
-/* XXX peek_u8 / peek_u32 */
 
 /*
  * Functions to extract or store SSH wire encoded bignums and elliptic