From 7283b99bb75ab8e1d6d3c7d02011bfb6bd318305 Mon Sep 17 00:00:00 2001
From: Alexander Andryashin <aandryashin@gmail.com>
Date: Tue, 1 Feb 2022 17:34:43 +0300
Subject: [PATCH] Add firstMatch capabilities.

---
 pprof.go |  1 +
 proxy.go | 24 ++++++++++++++++++------
 watch.go |  1 +
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/pprof.go b/pprof.go
index 4526fa1..773ea04 100644
--- a/pprof.go
+++ b/pprof.go
@@ -1,3 +1,4 @@
+//go:build pprof
 // +build pprof
 
 package main
diff --git a/proxy.go b/proxy.go
index f21f4a0..2cc79e0 100644
--- a/proxy.go
+++ b/proxy.go
@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"encoding/xml"
 	"fmt"
-	"github.com/abbot/go-http-auth"
 	"io"
 	"io/ioutil"
 	"log"
@@ -20,6 +19,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/abbot/go-http-auth"
 	. "github.com/aerokube/ggr/config"
 	"golang.org/x/net/websocket"
 )
@@ -61,10 +61,12 @@ var keys = struct {
 	desiredCapabilities string
 	w3cCapabilities     string
 	alwaysMatch         string
+	firstMatch          string
 }{
 	desiredCapabilities: "desiredCapabilities",
 	w3cCapabilities:     "capabilities",
 	alwaysMatch:         "alwaysMatch",
+	firstMatch:          "firstMatch",
 }
 
 var (
@@ -95,13 +97,23 @@ func (c caps) capabilities(fn func(m map[string]interface{}, w3c bool, extension
 	}
 	if w3cCapabilities, ok := c[keys.w3cCapabilities]; ok {
 		if m, ok := w3cCapabilities.(map[string]interface{}); ok {
+			var match map[string]interface{}
 			if alwaysMatch, ok := m[keys.alwaysMatch]; ok {
 				if m, ok := alwaysMatch.(map[string]interface{}); ok {
-					fn(m, true, false)
-					for k, v := range m { // Extension capabilities have ":" in key
-						if ec, ok := v.(map[string]interface{}); ok && strings.Contains(k, ":") {
-							fn(ec, true, true)
-						}
+					match = m
+				}
+			} else if firstMatch, ok := m[keys.firstMatch]; ok {
+				if m, ok := firstMatch.([]interface{}); ok && len(m) > 0 {
+					if m, ok := m[0].(map[string]interface{}); ok {
+						match = m
+					}
+				}
+			}
+			if match != nil {
+				fn(match, true, false)
+				for k, v := range m { // Extension capabilities have ":" in key
+					if ec, ok := v.(map[string]interface{}); ok && strings.Contains(k, ":") {
+						fn(ec, true, true)
 					}
 				}
 			}
diff --git a/watch.go b/watch.go
index 438b48b..c912ce0 100644
--- a/watch.go
+++ b/watch.go
@@ -1,3 +1,4 @@
+//go:build watch
 // +build watch
 
 package main