Add demuxer probing

This allows auto-detection between section5 and annexb files, which
share the same extension.

6 files changed, 222 insertions, 28 deletions

diff --git a/tools/input/annexb.c b/tools/input/annexb.c
index a36536b..d99b1b0 100644
--- a/tools/input/annexb.c
+++ b/tools/input/annexb.c
@@ -1,6 +1,7 @@
 /*
  * Copyright © 2018, VideoLAN and dav1d authors
  * Copyright © 2018, Two Orioles, LLC
+ * Copyright © 2019, James Almer <jamrial@gmail.com>
  * All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
@@ -32,9 +33,92 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "common/intops.h"
+
+#include "dav1d/headers.h"
+
 #include "input/demuxer.h"
 #include "input/parse.h"
 
+// these functions are based on an implementation from FFmpeg, and relicensed
+// with author's permission
+
+#define PROBE_SIZE 1024
+
+static int annexb_probe(const uint8_t *data) {
+    int ret, cnt = 0;
+
+    size_t temporal_unit_size;
+    ret = leb(data + cnt, PROBE_SIZE - cnt, &temporal_unit_size);
+    if (ret < 0)
+        return 0;
+    cnt += ret;
+
+    size_t frame_unit_size;
+    ret = leb(data + cnt, PROBE_SIZE - cnt, &frame_unit_size);
+    if (ret < 0 || ((uint64_t)frame_unit_size + ret) > temporal_unit_size)
+        return 0;
+    cnt += ret;
+
+    temporal_unit_size -= ret;
+
+    size_t obu_unit_size;
+    ret = leb(data + cnt, PROBE_SIZE - cnt, &obu_unit_size);
+    if (ret < 0 || ((uint64_t)obu_unit_size + ret) >= frame_unit_size)
+        return 0;
+    cnt += ret;
+
+    temporal_unit_size -= obu_unit_size + ret;
+    frame_unit_size -= obu_unit_size + ret;
+
+    // Check that the first OBU is a Temporal Delimiter.
+    size_t obu_size;
+    enum Dav1dObuType type;
+    ret = parse_obu_header(data + cnt, imin(PROBE_SIZE - cnt, (int) obu_unit_size),
+                           &obu_size, &type, 1);
+    if (ret < 0 || type != DAV1D_OBU_TD || obu_size > 0)
+        return 0;
+    cnt += obu_unit_size;
+
+    // look for first frame and accompanying sequence header
+    int seq = 0;
+    while (cnt < PROBE_SIZE) {
+        ret = leb(data + cnt, PROBE_SIZE - cnt, &obu_unit_size);
+        if (ret < 0 || ((uint64_t)obu_unit_size + ret) > frame_unit_size)
+            return 0;
+        cnt += ret;
+        temporal_unit_size -= ret;
+        frame_unit_size -= ret;
+
+        ret = parse_obu_header(data + cnt, imin(PROBE_SIZE - cnt, (int) obu_unit_size),
+                               &obu_size, &type, 1);
+        if (ret < 0)
+            return 0;
+        cnt += obu_unit_size;
+
+        switch (type) {
+        case DAV1D_OBU_SEQ_HDR:
+            seq = 1;
+            break;
+        case DAV1D_OBU_FRAME:
+        case DAV1D_OBU_FRAME_HDR:
+            return seq;
+        case DAV1D_OBU_TD:
+        case DAV1D_OBU_TILE_GRP:
+            return 0;
+        default:
+            break;
+        }
+
+        temporal_unit_size -= obu_unit_size;
+        frame_unit_size -= obu_unit_size;
+        if (frame_unit_size <= 0)
+            break;
+    }
+
+    return 0;
+}
+
 typedef struct DemuxerPriv {
     FILE *f;
     size_t temporal_unit_size;
@@ -103,7 +187,8 @@ static void annexb_close(AnnexbInputContext *const c) {
 const Demuxer annexb_demuxer = {
     .priv_data_size = sizeof(AnnexbInputContext),
     .name = "annexb",
-    .extension = "obu",
+    .probe = annexb_probe,
+    .probe_sz = PROBE_SIZE,
     .open = annexb_open,
     .read = annexb_read,
     .close = annexb_close,
diff --git a/tools/input/demuxer.h b/tools/input/demuxer.h
index ef79752..c2b88e1 100644
--- a/tools/input/demuxer.h
+++ b/tools/input/demuxer.h
@@ -34,7 +34,8 @@ typedef struct DemuxerPriv DemuxerPriv;
 typedef struct Demuxer {
     int priv_data_size;
     const char *name;
-    const char *extension;
+    int probe_sz;
+    int (*probe)(const uint8_t *data);
     int (*open)(DemuxerPriv *ctx, const char *filename,
                 unsigned fps[2], unsigned *num_frames, unsigned timebase[2]);
     int (*read)(DemuxerPriv *ctx, Dav1dData *data);
diff --git a/tools/input/input.c b/tools/input/input.c
index 5dc4cef..ec5a418 100644
--- a/tools/input/input.c
+++ b/tools/input/input.c
@@ -33,6 +33,7 @@
 #include <string.h>
 
 #include "common/attributes.h"
+#include "common/intops.h"
 
 #include "input/input.h"
 #include "input/demuxer.h"
@@ -58,23 +59,6 @@ void init_demuxers(void) {
     register_demuxer(section5_demuxer);
 }
 
-static const char *find_extension(const char *const f) {
-    const size_t l = strlen(f);
-
-    if (l == 0) return NULL;
-
-    const char *const end = &f[l - 1], *step = end;
-    while ((*step >= 'a' && *step <= 'z') ||
-           (*step >= 'A' && *step <= 'Z') ||
-           (*step >= '0' && *step <= '9'))
-    {
-        step--;
-    }
-
-    return (step < end && step > f && *step == '.' && step[-1] != '/') ?
-           &step[1] : NULL;
-}
-
 int input_open(DemuxerContext **const c_out,
                const char *const name, const char *const filename,
                unsigned fps[2], unsigned *const num_frames, unsigned timebase[2])
@@ -95,22 +79,34 @@ int input_open(DemuxerContext **const c_out,
             return DAV1D_ERR(ENOPROTOOPT);
         }
     } else {
-        const char *const ext = find_extension(filename);
-        if (!ext) {
-            fprintf(stderr, "No extension found for file %s\n", filename);
-            return -1;
+        int probe_sz = 0;
+        for (i = 0; i < num_demuxers; i++)
+            probe_sz = imax(probe_sz, demuxers[i]->probe_sz);
+        uint8_t *const probe_data = malloc(probe_sz);
+        if (!probe_data) {
+            fprintf(stderr, "Failed to allocate memory\n");
+            return DAV1D_ERR(ENOMEM);
+        }
+        FILE *f = fopen(filename, "rb");
+        res = !!fread(probe_data, 1, probe_sz, f);
+        fclose(f);
+        if (!res) {
+            free(probe_data);
+            fprintf(stderr, "Failed to read probe data\n");
+            return errno ? DAV1D_ERR(errno) : DAV1D_ERR(ENODATA);
         }
 
         for (i = 0; i < num_demuxers; i++) {
-            if (!strcmp(demuxers[i]->extension, ext)) {
+            if (demuxers[i]->probe(probe_data)) {
                 impl = demuxers[i];
                 break;
             }
         }
+        free(probe_data);
         if (i == num_demuxers) {
             fprintf(stderr,
-                    "Failed to find demuxer for file %s (\"%s\")\n",
-                    filename, ext);
+                    "Failed to probe demuxer for file %s\n",
+                    filename);
             return DAV1D_ERR(ENOPROTOOPT);
         }
     }
diff --git a/tools/input/ivf.c b/tools/input/ivf.c
index e83dd6d..15676b0 100644
--- a/tools/input/ivf.c
+++ b/tools/input/ivf.c
@@ -39,6 +39,16 @@ typedef struct DemuxerPriv {
     FILE *f;
 } IvfInputContext;
 
+static const uint8_t probe_data[] = {
+    'D', 'K', 'I', 'F',
+    0, 0, 0x20, 0,
+    'A', 'V', '0', '1',
+};
+
+static int ivf_probe(const uint8_t *const data) {
+    return !memcmp(data, probe_data, sizeof(probe_data));
+}
+
 static unsigned rl32(const uint8_t *const p) {
     return ((uint32_t)p[3] << 24U) | (p[2] << 16U) | (p[1] << 8U) | p[0];
 }
@@ -121,7 +131,8 @@ static void ivf_close(IvfInputContext *const c) {
 const Demuxer ivf_demuxer = {
     .priv_data_size = sizeof(IvfInputContext),
     .name = "ivf",
-    .extension = "ivf",
+    .probe = ivf_probe,
+    .probe_sz = sizeof(probe_data),
     .open = ivf_open,
     .read = ivf_read,
     .close = ivf_close,
diff --git a/tools/input/parse.h b/tools/input/parse.h
index 6ed3347..bebea21 100644
--- a/tools/input/parse.h
+++ b/tools/input/parse.h
@@ -48,4 +48,60 @@ static int leb128(FILE *const f, size_t *const len) {
     return i;
 }
 
+// these functions are based on an implementation from FFmpeg, and relicensed
+// with author's permission
+
+static int leb(const uint8_t *ptr, int sz, size_t *const len) {
+    unsigned i = 0, more;
+    *len = 0;
+    do {
+        if (!sz--) return -1;
+        const int byte = *ptr++;
+        more = byte & 0x80;
+        const unsigned bits = byte & 0x7f;
+        if (i <= 3 || (i == 4 && bits < (1 << 4)))
+            *len |= bits << (i * 7);
+        else if (bits) return -1;
+        if (++i == 8 && more) return -1;
+    } while (more);
+    return i;
+}
+
+static inline int parse_obu_header(const uint8_t *buf, int buf_size,
+                                   size_t *const obu_size,
+                                   enum Dav1dObuType *const type,
+                                   const int allow_implicit_size)
+{
+    int ret, extension_flag, has_size_flag;
+
+    if (!buf_size)
+        return -1;
+    if (*buf & 0x80) // obu_forbidden_bit
+        return -1;
+
+    *type = (*buf & 0x78) >> 3;
+    extension_flag = (*buf & 0x4) >> 2;
+    has_size_flag  = (*buf & 0x2) >> 1;
+    // ignore obu_reserved_1bit
+    buf++;
+    buf_size--;
+
+    if (extension_flag) {
+        buf++;
+        buf_size--;
+        // ignore fields
+    }
+
+    if (has_size_flag) {
+        ret = leb(buf, buf_size, obu_size);
+        if (ret < 0)
+            return -1;
+        return (int) *obu_size + ret + 1 + extension_flag;
+    } else if (!allow_implicit_size)
+        return -1;
+
+    *obu_size = buf_size;
+    return buf_size + 1 + extension_flag;
+}
+
 #endif /* DAV1D_INPUT_PARSE_H */
diff --git a/tools/input/section5.c b/tools/input/section5.c
index 7df49a9..dfe5ab9 100644
--- a/tools/input/section5.c
+++ b/tools/input/section5.c
@@ -1,6 +1,7 @@
 /*
  * Copyright © 2019, VideoLAN and dav1d authors
  * Copyright © 2019, Two Orioles, LLC
+ * Copyright © 2019, James Almer <jamrial@gmail.com>
  * All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
@@ -32,9 +33,52 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "dav1d/headers.h"
+
 #include "input/demuxer.h"
 #include "input/parse.h"
 
+#define PROBE_SIZE 1024
+
+static int section5_probe(const uint8_t *data) {
+    int ret, cnt = 0;
+
+    // Check that the first OBU is a Temporal Delimiter.
+    size_t obu_size;
+    enum Dav1dObuType type;
+    ret = parse_obu_header(data + cnt, PROBE_SIZE - cnt,
+                           &obu_size, &type, 0);
+    if (ret < 0 || type != DAV1D_OBU_TD || obu_size > 0)
+        return 0;
+    cnt += ret;
+
+    // look for first frame and accompanying sequence header
+    int seq = 0;
+    while (cnt < PROBE_SIZE) {
+        ret = parse_obu_header(data + cnt, PROBE_SIZE - cnt,
+                               &obu_size, &type, 0);
+        if (ret < 0)
+            return 0;
+        cnt += ret;
+
+        switch (type) {
+        case DAV1D_OBU_SEQ_HDR:
+            seq = 1;
+            break;
+        case DAV1D_OBU_FRAME:
+        case DAV1D_OBU_FRAME_HDR:
+            return seq;
+        case DAV1D_OBU_TD:
+        case DAV1D_OBU_TILE_GRP:
+            return 0;
+        default:
+            break;
+        }
+    }
+
+    return 0;
+}
+
 typedef struct DemuxerPriv {
     FILE *f;
 } Section5InputContext;
@@ -132,7 +176,8 @@ static void section5_close(Section5InputContext *const c) {
 const Demuxer section5_demuxer = {
     .priv_data_size = sizeof(Section5InputContext),
     .name = "section5",
-    .extension = "obu",
+    .probe = section5_probe,
+    .probe_sz = PROBE_SIZE,
     .open = section5_open,
     .read = section5_read,
     .close = section5_close,