diff --git a/pkg/generate/config.go b/pkg/generate/config.go index e4b1664..c2bffcb 100644 --- a/pkg/generate/config.go +++ b/pkg/generate/config.go @@ -155,8 +155,10 @@ func newConfig(info *cameraInfo, req *Request) string { b.WriteString("go2rtc:\n streams:\n") writeStreamLines(&b, info) - writeCredentials(&b, info.Credentials) + if len(info.Credentials) == 0 { + b.WriteByte('\n') + } b.WriteString("cameras:\n") writeCameraBlock(&b, info, req) diff --git a/pkg/generate/insert.go b/pkg/generate/insert.go index 7a3d1c8..67a1591 100644 --- a/pkg/generate/insert.go +++ b/pkg/generate/insert.go @@ -15,6 +15,7 @@ var ( reStreamName = regexp.MustCompile(`^\s{4}'?(\w[\w-]*)'?:`) reStreamContent = regexp.MustCompile(`^\s{4,}`) reNextSection = regexp.MustCompile(`^[a-z#]`) + reSibling = regexp.MustCompile(`^ \w`) // sibling under go2rtc: (ex. ` xiaomi:`) reCameraBody = regexp.MustCompile(`^\s{2,}\S`) reVersion = regexp.MustCompile(`^version:`) ) @@ -138,6 +139,13 @@ func findStreamInsertPoint(lines []string) int { continue } if in { + // stop at sibling under go2rtc: (ex. ` xiaomi:`) + if reSibling.MatchString(line) && !reStreamsHeader.MatchString(line) { + if last >= 0 { + return last + 1 + } + return headerIdx + 1 + } if reStreamContent.MatchString(line) { last = i } else if reNextSection.MatchString(line) { @@ -160,10 +168,10 @@ func findStreamInsertPoint(lines []string) int { return -1 } -// upsertCredentials merges creds into existing top-level sections. For each -// section: if a matching line ` "":` exists -- replace its value; else -// insert in sorted order. If the section itself doesn't exist -- create a new -// top-level block just before `cameras:`. +// upsertCredentials merges creds into credential sections nested under go2rtc:. +// For each section: if a matching line ` "":` exists -- replace its +// value; else insert in sorted order. If the section itself doesn't exist -- +// create a new nested block inside go2rtc: (after streams: block). func upsertCredentials(lines []string, creds map[string]map[string]string, added map[int]bool) ([]string, map[int]bool) { if len(creds) == 0 { return lines, added @@ -182,11 +190,12 @@ func upsertCredentials(lines []string, creds map[string]map[string]string, added return lines, added } -// upsertSection updates or appends a single top-level section. -var reCredKey = regexp.MustCompile(`^\s{2}"([^"]+)":`) +// ex. ` "4161148305": V1:xxx` -- 4-space indent under nested section +var reCredKey = regexp.MustCompile(`^\s{4}"([^"]+)":`) func upsertSection(lines []string, section string, kv map[string]string, added map[int]bool) ([]string, map[int]bool) { - reHeader := regexp.MustCompile(`^` + regexp.QuoteMeta(section) + `:\s*$`) + // section header is nested under go2rtc:, ex. ` xiaomi:` + reHeader := regexp.MustCompile(`^ ` + regexp.QuoteMeta(section) + `:\s*$`) headerIdx := -1 for i, line := range lines { @@ -200,10 +209,16 @@ func upsertSection(lines []string, section string, kv map[string]string, added m return insertNewSection(lines, section, kv, added) } - // section exists -- find last content line of the section (skip trailing blanks) + // section exists -- find end (blank line, top-level header, or sibling 2-space key) end := len(lines) for i := headerIdx + 1; i < len(lines); i++ { - if strings.TrimSpace(lines[i]) == "" || reTopLevel.MatchString(lines[i]) { + line := lines[i] + if strings.TrimSpace(line) == "" || reTopLevel.MatchString(line) { + end = i + break + } + // sibling under go2rtc: has 2-space indent, not 4 + if len(line) >= 2 && line[0] == ' ' && line[1] == ' ' && (len(line) == 2 || line[2] != ' ') { end = i break } @@ -216,9 +231,9 @@ func upsertSection(lines []string, section string, kv map[string]string, added m sort.Strings(keys) for _, k := range keys { - newLine := fmt.Sprintf(" %q: %s", k, kv[k]) + newLine := fmt.Sprintf(" %q: %s", k, kv[k]) - // try replace -- no length change, just mark modified line as added + // try replace -- no length change, just mark modified line replaced := false for i := headerIdx + 1; i < end; i++ { if m := reCredKey.FindStringSubmatch(lines[i]); m != nil && m[1] == k { @@ -257,35 +272,24 @@ func upsertSection(lines []string, section string, kv map[string]string, added m return lines, added } +// insertNewSection adds a new nested section under go2rtc:, after the streams: +// block but before any sibling go2rtc key or top-level header. func insertNewSection(lines []string, section string, kv map[string]string, added map[int]bool) ([]string, map[int]bool) { - camIdx := -1 - for i, line := range lines { - if reCamerasHeader.MatchString(line) { - camIdx = i - break - } - } - if camIdx == -1 { - camIdx = len(lines) + // find end of streams: block inside go2rtc: + insertAt := findGo2RTCInsertPoint(lines) + if insertAt < 0 { + return lines, added } - // insert point: right before the blank line that precedes cameras: - // keep one blank line between blocks - insertAt := camIdx - for insertAt > 0 && strings.TrimSpace(lines[insertAt-1]) == "" { - insertAt-- - } - - block := []string{section + ":"} + block := []string{" " + section + ":"} keys := make([]string, 0, len(kv)) for k := range kv { keys = append(keys, k) } sort.Strings(keys) for _, k := range keys { - block = append(block, fmt.Sprintf(" %q: %s", k, kv[k])) + block = append(block, fmt.Sprintf(" %q: %s", k, kv[k])) } - block = append(block, "") lines = append(lines[:insertAt], append(block, lines[insertAt:]...)...) added = shiftAdded(added, insertAt) @@ -296,6 +300,38 @@ func insertNewSection(lines []string, section string, kv map[string]string, adde return lines, added } +// findGo2RTCInsertPoint returns the line index where a new nested section +// under go2rtc: should be inserted -- after the last non-blank content line +// of the go2rtc: block. +func findGo2RTCInsertPoint(lines []string) int { + reGo2RTCHeader := regexp.MustCompile(`^go2rtc:\s*$`) + + headerIdx := -1 + for i, line := range lines { + if reGo2RTCHeader.MatchString(line) { + headerIdx = i + break + } + } + if headerIdx == -1 { + return -1 + } + + last := headerIdx + for i := headerIdx + 1; i < len(lines); i++ { + line := lines[i] + if strings.TrimSpace(line) == "" { + continue + } + if reTopLevel.MatchString(line) { + break + } + last = i + } + + return last + 1 +} + // shiftAdded moves all marks at index >= from by +1. Also used with from=len(lines) // as a no-op shift (just return same map). func shiftAdded(added map[int]bool, from int) map[int]bool { diff --git a/pkg/generate/writer.go b/pkg/generate/writer.go index dfb830d..6fde79a 100644 --- a/pkg/generate/writer.go +++ b/pkg/generate/writer.go @@ -14,15 +14,17 @@ func writeStreamLines(b *strings.Builder, info *cameraInfo) { fmt.Fprintf(b, " '%s':\n", info.SubStreamName) fmt.Fprintf(b, " - %s\n", info.SubSource) } - - b.WriteByte('\n') } -// writeCredentials writes top-level credential sections (xiaomi, tapo, ring, ...) -// populated by registered ExtractFunc handlers. Sorted by section, then by key. +// writeCredentials writes credential sections (xiaomi, tapo, ring, ...) as +// nested keys under go2rtc:, populated by registered ExtractFunc handlers. +// Sorted by section, then by key. Frigate only allows known keys at root, +// so credentials must live inside go2rtc: (which allows extra keys). // ex. -// xiaomi: -// "4161148305": V1:9d2w... +// go2rtc: +// streams: { ... } +// xiaomi: +// "4161148305": V1:9d2w... func writeCredentials(b *strings.Builder, creds map[string]map[string]string) { if len(creds) == 0 { return @@ -35,7 +37,7 @@ func writeCredentials(b *strings.Builder, creds map[string]map[string]string) { sort.Strings(sections) for _, section := range sections { - fmt.Fprintf(b, "%s:\n", section) + fmt.Fprintf(b, " %s:\n", section) keys := make([]string, 0, len(creds[section])) for k := range creds[section] { @@ -44,10 +46,10 @@ func writeCredentials(b *strings.Builder, creds map[string]map[string]string) { sort.Strings(keys) for _, k := range keys { - fmt.Fprintf(b, " %q: %s\n", k, creds[section][k]) + fmt.Fprintf(b, " %q: %s\n", k, creds[section][k]) } - b.WriteByte('\n') } + b.WriteByte('\n') } func writeCameraBlock(b *strings.Builder, info *cameraInfo, req *Request) { diff --git a/pkg/generate/xiaomi_test.go b/pkg/generate/xiaomi_test.go new file mode 100644 index 0000000..f7477d5 --- /dev/null +++ b/pkg/generate/xiaomi_test.go @@ -0,0 +1,396 @@ +package generate + +import ( + "net/url" + "strings" + "sync" + "testing" +) + +// registerXiaomi installs a xiaomi extractor identical to the one in +// internal/xiaomi. Tests live here (not in internal/xiaomi) because they +// validate generator behavior with xiaomi-style URLs. +var registerOnce sync.Once + +func registerXiaomi() { + registerOnce.Do(func() { + RegisterExtract("xiaomi", func(rawURL string) (cleaned, section, key, value string) { + u, err := url.Parse(rawURL) + if err != nil || u.User == nil { + return rawURL, "", "", "" + } + q := u.Query() + token := q.Get("token") + if token == "" { + return rawURL, "", "", "" + } + q.Del("token") + u.RawQuery = q.Encode() + return u.String(), "xiaomi", u.User.Username(), token + }) + }) +} + +// --- Helpers --- + +func mustGen(t *testing.T, req *Request) string { + t.Helper() + r, err := Generate(req) + if err != nil { + t.Fatalf("Generate: %v", err) + } + return r.Config +} + +func assertContains(t *testing.T, cfg, substr string) { + t.Helper() + if !strings.Contains(cfg, substr) { + t.Errorf("expected config to contain:\n %q\n--- got ---\n%s", substr, cfg) + } +} + +func assertNotContains(t *testing.T, cfg, substr string) { + t.Helper() + if strings.Contains(cfg, substr) { + t.Errorf("expected config NOT to contain:\n %q\n--- got ---\n%s", substr, cfg) + } +} + +func countOccurrences(s, substr string) int { + if substr == "" { + return 0 + } + n := 0 + for i := 0; ; { + j := strings.Index(s[i:], substr) + if j < 0 { + return n + } + n++ + i += j + len(substr) + } +} + +// ex. "xiaomi://user:cn@ip?did=D&model=M&token=T" +func xurl(user, region, ip, did, model, token string) string { + return "xiaomi://" + user + ":" + region + "@" + ip + + "?did=" + did + "&model=" + model + "&token=" + url.QueryEscape(token) +} + +// --- Tests --- + +// Single xiaomi camera in a fresh config. +func TestXiaomi_NewConfig_SingleCamera(t *testing.T) { + registerXiaomi() + + cfg := mustGen(t, &Request{ + MainStream: xurl("acc1", "cn", "10.0.20.229", "1", "chuangmi.camera.v1", "V1:TOK_A"), + }) + + assertContains(t, cfg, "go2rtc:\n streams:\n") + assertContains(t, cfg, " '10_0_20_229_main':\n") + assertContains(t, cfg, "- xiaomi://acc1:cn@10.0.20.229?did=1&model=chuangmi.camera.v1\n") + assertContains(t, cfg, " xiaomi:\n \"acc1\": V1:TOK_A\n") + assertNotContains(t, cfg, "token=") + assertNotContains(t, cfg, "\nxiaomi:") // must be nested, not top-level +} + +// Two cameras on the same account -- token appears only once. +func TestXiaomi_SameAccount_TokenNotDuplicated(t *testing.T) { + registerXiaomi() + + c1 := mustGen(t, &Request{ + MainStream: xurl("acc1", "cn", "10.0.20.229", "1", "v1", "V1:TOK_A"), + }) + c2 := mustGen(t, &Request{ + MainStream: xurl("acc1", "cn", "10.0.20.230", "2", "v2", "V1:TOK_A"), + ExistingConfig: c1, + }) + + if n := countOccurrences(c2, `"acc1":`); n != 1 { + t.Errorf("expected exactly 1 \"acc1\" key, got %d\n---\n%s", n, c2) + } + assertContains(t, c2, `"acc1": V1:TOK_A`) + assertContains(t, c2, " '10_0_20_229_main':") + assertContains(t, c2, " '10_0_20_230_main':") +} + +// Two accounts -- both tokens present, sorted by key. +func TestXiaomi_TwoAccounts_SortedKeys(t *testing.T) { + registerXiaomi() + + c1 := mustGen(t, &Request{ + MainStream: xurl("zeta", "cn", "10.0.20.229", "1", "v1", "TOK_Z"), + }) + c2 := mustGen(t, &Request{ + MainStream: xurl("alpha", "de", "10.0.20.230", "2", "v2", "TOK_A"), + ExistingConfig: c1, + }) + + iAlpha := strings.Index(c2, `"alpha":`) + iZeta := strings.Index(c2, `"zeta":`) + if iAlpha < 0 || iZeta < 0 { + t.Fatalf("expected both keys:\n%s", c2) + } + if iAlpha >= iZeta { + t.Errorf("expected alpha before zeta (sorted)\n%s", c2) + } +} + +// Re-login with a new token overwrites the existing value. +func TestXiaomi_TokenRefresh_OverwritesValue(t *testing.T) { + registerXiaomi() + + c1 := mustGen(t, &Request{ + MainStream: xurl("acc1", "cn", "10.0.20.229", "1", "v1", "V1:OLD"), + }) + c2 := mustGen(t, &Request{ + MainStream: xurl("acc1", "cn", "10.0.20.230", "2", "v2", "V1:NEW"), + ExistingConfig: c1, + }) + + assertContains(t, c2, `"acc1": V1:NEW`) + assertNotContains(t, c2, `"acc1": V1:OLD`) + if n := countOccurrences(c2, `"acc1":`); n != 1 { + t.Errorf("expected 1 key after refresh, got %d", n) + } +} + +// Main + Sub stream with same credentials -- token deduped to one entry. +func TestXiaomi_MainAndSub_SameAccount_OneToken(t *testing.T) { + registerXiaomi() + + cfg := mustGen(t, &Request{ + MainStream: xurl("acc1", "cn", "10.0.20.229", "1", "main", "V1:TOK"), + SubStream: xurl("acc1", "cn", "10.0.20.229", "1", "sub", "V1:TOK"), + }) + + if n := countOccurrences(cfg, `"acc1":`); n != 1 { + t.Errorf("expected 1 acc1 key for main+sub, got %d\n%s", n, cfg) + } + assertContains(t, cfg, " '10_0_20_229_main':") + assertContains(t, cfg, " '10_0_20_229_sub':") + assertNotContains(t, cfg, "token=") +} + +// Main and Sub from different accounts -- both tokens in the section. +func TestXiaomi_MainAndSub_DifferentAccounts(t *testing.T) { + registerXiaomi() + + cfg := mustGen(t, &Request{ + MainStream: xurl("accA", "cn", "10.0.20.229", "1", "v1", "TOK_A"), + SubStream: xurl("accB", "de", "10.0.20.229", "1", "v1", "TOK_B"), + }) + + assertContains(t, cfg, `"accA": TOK_A`) + assertContains(t, cfg, `"accB": TOK_B`) +} + +// 10 cameras across 3 accounts added sequentially -- exactly 3 tokens at the end, +// correct token values, all streams present. +func TestXiaomi_Scale_10Cameras_3Accounts(t *testing.T) { + registerXiaomi() + + cases := []struct{ user, ip, token string }{ + {"accA", "10.0.20.10", "TOK_A_v1"}, + {"accA", "10.0.20.11", "TOK_A_v1"}, + {"accB", "10.0.20.12", "TOK_B_v1"}, + {"accA", "10.0.20.13", "TOK_A_v1"}, + {"accC", "10.0.20.14", "TOK_C_v1"}, + {"accB", "10.0.20.15", "TOK_B_v2"}, // B gets refreshed + {"accC", "10.0.20.16", "TOK_C_v1"}, + {"accA", "10.0.20.17", "TOK_A_v2"}, // A gets refreshed + {"accB", "10.0.20.18", "TOK_B_v2"}, + {"accC", "10.0.20.19", "TOK_C_v2"}, // C gets refreshed + } + + cfg := "" + for i, c := range cases { + req := &Request{ + MainStream: xurl(c.user, "cn", c.ip, "1", "v", c.token), + } + if i > 0 { + req.ExistingConfig = cfg + } + cfg = mustGen(t, req) + } + + if n := countOccurrences(cfg, `"accA":`); n != 1 { + t.Errorf("accA: expected 1 key, got %d", n) + } + if n := countOccurrences(cfg, `"accB":`); n != 1 { + t.Errorf("accB: expected 1 key, got %d", n) + } + if n := countOccurrences(cfg, `"accC":`); n != 1 { + t.Errorf("accC: expected 1 key, got %d", n) + } + + // final (latest) tokens + assertContains(t, cfg, `"accA": TOK_A_v2`) + assertContains(t, cfg, `"accB": TOK_B_v2`) + assertContains(t, cfg, `"accC": TOK_C_v2`) + + for _, c := range cases { + want := "xiaomi://" + c.user + ":cn@" + c.ip + if !strings.Contains(cfg, want) { + t.Errorf("missing stream URL %q", want) + } + } + + // only one xiaomi: section header, only one go2rtc: + if n := countOccurrences(cfg, "\n xiaomi:\n"); n != 1 { + t.Errorf("expected 1 xiaomi: header, got %d", n) + } + if n := countOccurrences(cfg, "\ngo2rtc:\n"); n != 1 { + t.Errorf("expected 1 go2rtc: header, got %d", n) + } +} + +// URL without ?token=... -- extractor returns empty section, no xiaomi: block written. +func TestXiaomi_URLWithoutToken_NoSection(t *testing.T) { + registerXiaomi() + + cfg := mustGen(t, &Request{ + MainStream: "xiaomi://acc1:cn@10.0.20.229?did=1&model=v1", + }) + + // the nested section header would look like "\n xiaomi:\n" -- URL scheme + // "xiaomi://" must not trigger a false positive + assertNotContains(t, cfg, "\n xiaomi:\n") + assertContains(t, cfg, "- xiaomi://acc1:cn@10.0.20.229?did=1&model=v1\n") +} + +// Malformed URL must not crash the generator; URL is passed through as-is. +func TestXiaomi_MalformedURL_DoesNotPanic(t *testing.T) { + registerXiaomi() + + _, err := Generate(&Request{ + MainStream: "xiaomi://%%%bad", + }) + if err != nil { + t.Logf("Generate returned error (ok): %v", err) + } +} + +// Token with base64 special chars (+ / =) must survive YAML write without escaping. +func TestXiaomi_TokenSpecialChars_PreservedRaw(t *testing.T) { + registerXiaomi() + + raw := "V1:9d2w+abc/def=end=" + cfg := mustGen(t, &Request{ + MainStream: xurl("acc1", "cn", "10.0.20.229", "1", "v1", raw), + }) + + assertContains(t, cfg, `"acc1": `+raw) +} + +// Go2RTC override MainStreamSource must also pass through the extractor. +func TestXiaomi_Go2RTCOverride_PassesThroughExtractor(t *testing.T) { + registerXiaomi() + + cfg := mustGen(t, &Request{ + MainStream: "rtsp://placeholder:554/stream", + Go2RTC: &Go2RTCOverride{ + MainStreamSource: xurl("acc1", "cn", "10.0.20.229", "1", "v1", "V1:OVR"), + }, + }) + + assertContains(t, cfg, `"acc1": V1:OVR`) + assertNotContains(t, cfg, "token=") + assertContains(t, cfg, "- xiaomi://acc1:cn@10.0.20.229?did=1&model=v1\n") +} + +// addToConfig: existing config has no xiaomi: section -- must create one. +func TestXiaomi_AddToConfig_NoExistingSection(t *testing.T) { + registerXiaomi() + + // start from a rtsp-only config + c1 := mustGen(t, &Request{ + MainStream: "rtsp://user:pass@10.0.20.100/stream1", + }) + assertNotContains(t, c1, "xiaomi:") + + c2 := mustGen(t, &Request{ + MainStream: xurl("acc1", "cn", "10.0.20.229", "1", "v1", "V1:TOK"), + ExistingConfig: c1, + }) + + assertContains(t, c2, " xiaomi:\n \"acc1\": V1:TOK\n") + assertContains(t, c2, "- rtsp://user:pass@10.0.20.100/stream1") + assertContains(t, c2, "- xiaomi://acc1:cn@10.0.20.229?did=1&model=v1") +} + +// addToConfig: existing config already has xiaomi: section with other accounts. +func TestXiaomi_AddToConfig_ExistingSection(t *testing.T) { + registerXiaomi() + + c1 := mustGen(t, &Request{ + MainStream: xurl("accA", "cn", "10.0.20.10", "1", "v1", "TOK_A"), + }) + c2 := mustGen(t, &Request{ + MainStream: xurl("accB", "de", "10.0.20.20", "2", "v2", "TOK_B"), + ExistingConfig: c1, + }) + + assertContains(t, c2, `"accA": TOK_A`) + assertContains(t, c2, `"accB": TOK_B`) + // xiaomi: section stays nested, exactly one header + if n := countOccurrences(c2, "\n xiaomi:\n"); n != 1 { + t.Errorf("expected 1 xiaomi header, got %d\n%s", n, c2) + } +} + +// Stream and camera names stay clean (no leftover tokens in URLs) with custom Name. +func TestXiaomi_CustomName_URLStillClean(t *testing.T) { + registerXiaomi() + + cfg := mustGen(t, &Request{ + Name: "my_cam", + MainStream: xurl("acc1", "cn", "10.0.20.229", "1", "v1", "V1:TOK"), + }) + + assertContains(t, cfg, " 'my_cam_main':") + assertContains(t, cfg, " my_cam:") + assertNotContains(t, cfg, "token=") + assertContains(t, cfg, `"acc1": V1:TOK`) +} + +// Mixed: rtsp + xiaomi -- rtsp URL untouched, xiaomi token extracted. +func TestXiaomi_MixedProtocols(t *testing.T) { + registerXiaomi() + + c1 := mustGen(t, &Request{ + MainStream: "rtsp://admin:pw@10.0.20.100/Streaming/Channels/101", + }) + c2 := mustGen(t, &Request{ + MainStream: xurl("acc1", "cn", "10.0.20.229", "1", "v1", "V1:TOK"), + ExistingConfig: c1, + }) + + assertContains(t, c2, "- rtsp://admin:pw@10.0.20.100/Streaming/Channels/101") + assertContains(t, c2, "- xiaomi://acc1:cn@10.0.20.229?did=1&model=v1") + assertContains(t, c2, `"acc1": V1:TOK`) +} + +// Order in generated config: go2rtc -> (streams, xiaomi) -> cameras -> version. +func TestXiaomi_SectionOrder(t *testing.T) { + registerXiaomi() + + cfg := mustGen(t, &Request{ + MainStream: xurl("acc1", "cn", "10.0.20.229", "1", "v1", "V1:TOK"), + }) + + iGo2rtc := strings.Index(cfg, "\ngo2rtc:\n") + iStreams := strings.Index(cfg, " streams:") + iXiaomi := strings.Index(cfg, " xiaomi:") + iCameras := strings.Index(cfg, "\ncameras:\n") + iVersion := strings.Index(cfg, "\nversion:") + + if iGo2rtc < 0 || iStreams < 0 || iXiaomi < 0 || iCameras < 0 || iVersion < 0 { + t.Fatalf("missing section in config:\n%s", cfg) + } + if !(iGo2rtc < iStreams && iStreams < iXiaomi && iXiaomi < iCameras && iCameras < iVersion) { + t.Errorf("wrong section order: go2rtc=%d streams=%d xiaomi=%d cameras=%d version=%d\n%s", + iGo2rtc, iStreams, iXiaomi, iCameras, iVersion, cfg) + } +}