[WIP] "Only allow screenshots" filter

community/manager.go

@@ -135,8 +135,9 @@ func (m *Manager) getCommunityFilterSet(ctx context.Context, communityId string)
 		postfilters = append(postfilters, filter.HellbanPostfilterName)
 	}
 	if m.instanceConfig.OpenAIApiKey != "" {
-		// Access to this filter is gated by further instance config (namely, the room IDs allowed to use it)
+		// Access to these filters is gated by further instance config (namely, the room IDs allowed to use them)
 		filters = append(filters, filter.OpenAIFilterName)
+		filters = append(filters, filter.ScreenshotClassifierFilterName)
 	}
 	if !communityConfig.StickyEventsFilterAllowStickyEvents {
 		filters = append(filters, filter.StickyEventsFilterName)

filter/filter_screenshot_classifier.go

@@ -0,0 +1,168 @@
+package filter
+
+import (
+	"context"
+	"encoding/base64"
+	"fmt"
+	"log"
+	"net/http"
+	"slices"
+	"strings"
+	"sync"
+
+	"github.com/matrix-org/policyserv/config"
+	"github.com/matrix-org/policyserv/filter/classification"
+	"github.com/openai/openai-go/v3"
+	"github.com/openai/openai-go/v3/option"
+	"github.com/openai/openai-go/v3/responses"
+)
+
+type ScreenshotClassifier string
+
+const ScreenshotClassifierScreenshot ScreenshotClassifier = "screenshot"
+const ScreenshotClassifierPhoto ScreenshotClassifier = "photo"
+const ScreenshotClassifierUnknown ScreenshotClassifier = "unknown"
+const ScreenshotClassifierTooLarge ScreenshotClassifier = "too_large" // always causes a reject
+
+const ScreenshotClassifierFilterName = "ScreenshotClassifierFilter"
+
+func init() {
+	mustRegister(ScreenshotClassifierFilterName, &ScreenshotClassifierFilter{})
+}
+
+type ScreenshotClassifierFilter struct {
+}
+
+func (s *ScreenshotClassifierFilter) MakeFor(set *Set) (Instanced, error) {
+	return &InstancedScreenshotClassifierFilter{
+		set:        set,
+		aiProvider: NewOpenAIScreenshotClassifierAIProvider(set.instanceConfig),
+		inRoomIds:  set.instanceConfig.OpenAIAllowedRoomIds,
+	}, nil
+}
+
+type InstancedScreenshotClassifierFilter struct {
+	set        *Set
+	aiProvider AIProvider // we use indirection because we can't (realistically) test that OpenAI is working
+	inRoomIds  []string
+}
+
+func (f *InstancedScreenshotClassifierFilter) Name() string {
+	return ScreenshotClassifierFilterName
+}
+
+func (f *InstancedScreenshotClassifierFilter) CheckEvent(ctx context.Context, input *Input) ([]classification.Classification, error) {
+	//if !slices.Contains(f.inRoomIds, input.Event.RoomID().String()) {
+	//	return nil, nil
+	//}
+	return f.aiProvider.CheckEvent(ctx, &aiFilterConfig{
+		FailSecure: f.set.communityConfig.OpenAIFilterFailSecure,
+	}, input)
+}
+
+type OpenAIScreenshotClassifierAIProvider struct {
+	client openai.Client
+}
+
+func NewOpenAIScreenshotClassifierAIProvider(instanceConfig *config.InstanceConfig) *OpenAIScreenshotClassifierAIProvider {
+	client := openai.NewClient(option.WithAPIKey(instanceConfig.OpenAIApiKey))
+	return &OpenAIScreenshotClassifierAIProvider{
+		client: client,
+	}
+}
+
+func (p *OpenAIScreenshotClassifierAIProvider) CheckEvent(ctx context.Context, cnf *aiFilterConfig, input *Input) ([]classification.Classification, error) {
+	if len(input.Medias) == 0 {
+		return nil, nil // nothing to scan here
+	}
+
+	// Run the checks concurrently for all media items
+	mediaClassifications := make([]ScreenshotClassifier, len(input.Medias))
+	wg := sync.WaitGroup{}
+	for i, media := range input.Medias {
+		wg.Add(1)
+		go func(i int, media *Media) {
+			defer wg.Done()
+			var err error
+			mediaClassifications[i], err = p.classifyMedia(ctx, media)
+			if err != nil {
+				log.Printf("[%s | %s] Error classifying media %s: %s", input.Event.EventID(), input.Event.RoomID(), media.String(), err)
+				if cnf.FailSecure {
+					mediaClassifications[i] = ScreenshotClassifierTooLarge
+				} else {
+					mediaClassifications[i] = ScreenshotClassifierUnknown
+				}
+			}
+		}(i, media)
+	}
+
+	// Wait for all the work to complete
+	wg.Wait()
+
+	// Process the responses
+	log.Printf("[%s | %s] Screenshot classification results: %v", input.Event.EventID(), input.Event.RoomID(), mediaClassifications)
+	allowedTypes := []ScreenshotClassifier{ScreenshotClassifierScreenshot} // TODO: Config
+	isAllowed := false
+	for _, response := range mediaClassifications {
+		if response == ScreenshotClassifierTooLarge {
+			isAllowed = false
+			break // this is a security failure, so we don't want to continue
+		} else if slices.Contains(allowedTypes, response) {
+			isAllowed = true
+		}
+	}
+	if isAllowed {
+		return nil, nil
+	}
+	return []classification.Classification{classification.Spam}, nil
+}
+
+func (p *OpenAIScreenshotClassifierAIProvider) classifyMedia(ctx context.Context, media *Media) (ScreenshotClassifier, error) {
+	b, err := media.Download()
+	if err != nil {
+		return ScreenshotClassifierUnknown, fmt.Errorf("error downloading media: %w", err)
+	}
+	if len(b) > 50*1024*1024 { // 50mb
+		return ScreenshotClassifierTooLarge, nil
+	}
+
+	mimeType := http.DetectContentType(b)
+	if !strings.HasPrefix(mimeType, "image/") {
+		log.Printf("%s detected as %s (non-image), skipping", media.String(), mimeType)
+		return ScreenshotClassifierUnknown, nil
+	}
+
+	resp, err := p.client.Responses.New(ctx, responses.ResponseNewParams{
+		Model: "gpt-5-nano",
+		Input: responses.ResponseNewParamsInputUnion{
+			OfInputItemList: responses.ResponseInputParam{{
+				OfInputMessage: &responses.ResponseInputItemMessageParam{
+					Role: "user",
+					Content: responses.ResponseInputMessageContentListParam{{
+						OfInputText: &responses.ResponseInputTextParam{
+							Text: "Only respond with 'PHOTO', 'SCREENSHOT', or 'UNKNOWN' respectively. What is this?",
+						},
+					}, {
+						OfInputImage: &responses.ResponseInputImageParam{
+							Type:     "input_image",
+							ImageURL: openai.String(fmt.Sprintf("data:%s;base64,%s", mimeType, base64.StdEncoding.EncodeToString(b))),
+						},
+					}},
+				},
+			}},
+		},
+	})
+	if err != nil {
+		return ScreenshotClassifierUnknown, err
+	}
+
+	aiClass := strings.ToUpper(strings.TrimSpace(resp.OutputText()))
+	log.Printf("Classified %s as '%s'", media.String(), aiClass)
+	if strings.HasPrefix(aiClass, "PHOTO") {
+		return ScreenshotClassifierPhoto, nil
+	} else if strings.HasPrefix(aiClass, "SCREENSHOT") {
+		return ScreenshotClassifierScreenshot, nil
+	} else {
+		return ScreenshotClassifierUnknown, nil
+	}
+}

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/kelseyhightower/envconfig v1.4.0
 	github.com/lib/pq v1.10.9
 	github.com/matrix-org/gomatrixserverlib v0.0.0-20251103190139-eb14008fe6d1
+	github.com/openai/openai-go/v3 v3.11.0
 	github.com/panjf2000/ants/v2 v2.11.3
 	github.com/prometheus/client_golang v1.23.2
 	github.com/ryanuber/go-glob v1.0.0

go.sum

@@ -51,6 +51,8 @@ github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq
 github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
 github.com/oleiade/lane/v2 v2.0.0 h1:XW/ex/Inr+bPkLd3O240xrFOhUkTd4Wy176+Gv0E3Qw=
 github.com/oleiade/lane/v2 v2.0.0/go.mod h1:i5FBPFAYSWCgLh58UkUGCChjcCzef/MI7PlQm2TKCeg=
+github.com/openai/openai-go/v3 v3.11.0 h1:foJsKCgXqrx09WDyfI0beTuoIjTPFKBOx4/5jnw5o00=
+github.com/openai/openai-go/v3 v3.11.0/go.mod h1:cdufnVK14cWcT9qA1rRtrXx4FTRsgbDPW7Ia7SS5cZo=
 github.com/panjf2000/ants/v2 v2.11.3 h1:AfI0ngBoXJmYOpDh9m516vjqoUu2sLrIVgppI9TZVpg=
 github.com/panjf2000/ants/v2 v2.11.3/go.mod h1:8u92CYMUc6gyvTIw8Ru7Mt7+/ESnJahz5EVtqfrilek=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=