github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak/src/server/internal/broadcasts/broadcasts_test.go (about)

     1  // Copyright 2017 Google Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package broadcasts
    16  
    17  import (
    18  	"context"
    19  	"path"
    20  	"testing"
    21  	"time"
    22  
    23  	"github.com/google/fleetspeak/fleetspeak/src/common"
    24  	"github.com/google/fleetspeak/fleetspeak/src/comtesting"
    25  	"github.com/google/fleetspeak/fleetspeak/src/server/db"
    26  	"github.com/google/fleetspeak/fleetspeak/src/server/ids"
    27  	"github.com/google/fleetspeak/fleetspeak/src/server/internal/cache"
    28  	"github.com/google/fleetspeak/fleetspeak/src/server/internal/notifications"
    29  	"github.com/google/fleetspeak/fleetspeak/src/server/sqlite"
    30  
    31  	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
    32  	spb "github.com/google/fleetspeak/fleetspeak/src/server/proto/fleetspeak_server"
    33  	anypb "google.golang.org/protobuf/types/known/anypb"
    34  )
    35  
    36  func TestManager(t *testing.T) {
    37  	tmpDir, tmpDirCleanup := comtesting.GetTempDir("broadcasts")
    38  	defer tmpDirCleanup()
    39  	ds, err := sqlite.MakeDatastore(path.Join(tmpDir, "TestManager.sqlite"))
    40  	if err != nil {
    41  		t.Fatal(err)
    42  	}
    43  	ctx := context.Background()
    44  
    45  	var bid []ids.BroadcastID
    46  	for _, s := range []string{"0000000000000000", "0000000000000001", "0000000000000002"} {
    47  		b, err := ids.StringToBroadcastID(s)
    48  		if err != nil {
    49  			t.Fatalf("BroadcastID(%v) failed: %v", s, err)
    50  		}
    51  		bid = append(bid, b)
    52  	}
    53  
    54  	for i, tc := range []struct {
    55  		br  *spb.Broadcast
    56  		lim uint64
    57  	}{
    58  		{
    59  			br: &spb.Broadcast{
    60  				BroadcastId:    bid[0].Bytes(),
    61  				Source:         &fspb.Address{ServiceName: "testService"},
    62  				MessageType:    "WindowsBroadcast1",
    63  				RequiredLabels: []*fspb.Label{{ServiceName: "system", Label: "Windows"}},
    64  				Data: &anypb.Any{
    65  					TypeUrl: "message proto name 1",
    66  					Value:   []byte("message data 1"),
    67  				},
    68  			},
    69  			lim: 8,
    70  		},
    71  		{
    72  			br: &spb.Broadcast{
    73  				BroadcastId:    bid[1].Bytes(),
    74  				Source:         &fspb.Address{ServiceName: "testService"},
    75  				MessageType:    "LinuxBroadcast",
    76  				RequiredLabels: []*fspb.Label{{ServiceName: "system", Label: "Linux"}}},
    77  			lim: 8,
    78  		},
    79  	} {
    80  		if err := ds.CreateBroadcast(ctx, tc.br, tc.lim); err != nil {
    81  			t.Fatalf("%v: Unable to CreateBroadcast(%v): %v", i, tc.br, err)
    82  		}
    83  	}
    84  
    85  	bm, err := MakeManager(ctx, ds, 100*time.Millisecond, cache.NewClients(), notifications.NewDispatcher())
    86  	if err != nil {
    87  		t.Fatal(err)
    88  	}
    89  
    90  	clientID, err := common.BytesToClientID([]byte{0, 0, 0, 0, 0, 0, 0, 1})
    91  	if err != nil {
    92  		t.Fatal(err)
    93  	}
    94  
    95  	cd := &db.ClientData{
    96  		Key: []byte("A binary client key \x00\xff\x01\xfe"),
    97  		Labels: []*fspb.Label{
    98  			{ServiceName: "system", Label: "Windows"},
    99  			{ServiceName: "system", Label: "client-version-0.01"}},
   100  	}
   101  	if err := ds.AddClient(ctx, clientID, cd); err != nil {
   102  		t.Fatal(err)
   103  	}
   104  
   105  	msgs, err := bm.MakeBroadcastMessagesForClient(ctx, clientID, cd.Labels)
   106  	if err != nil {
   107  		t.Error(err)
   108  	} else {
   109  		if len(msgs) != 1 {
   110  			t.Errorf("Expected 1 broadcast messages, got: %v", msgs)
   111  		} else {
   112  			if msgs[0].MessageType != "WindowsBroadcast1" {
   113  				t.Errorf("Expected message of type [WindowsBroadcast1], got: [%v]", msgs[0])
   114  			}
   115  			if msgs[0].Data.TypeUrl != "message proto name 1" {
   116  				t.Errorf("Expected TypeUrl of [message proto name 1], got: [%v]", msgs[0].Data.TypeUrl)
   117  			}
   118  			if msgs[0].CreationTime == nil {
   119  				t.Error("Expected Creation time to be populated, was nil.")
   120  			}
   121  		}
   122  	}
   123  
   124  	msgs, err = bm.MakeBroadcastMessagesForClient(ctx, clientID, cd.Labels)
   125  	if err != nil {
   126  		t.Error(err)
   127  	} else {
   128  		if len(msgs) != 0 {
   129  			t.Errorf("Expected 0 broadcast messages on second run, got: %v", msgs)
   130  		}
   131  	}
   132  
   133  	br2 := &spb.Broadcast{
   134  		BroadcastId:    bid[2].Bytes(),
   135  		Source:         &fspb.Address{ServiceName: "testService"},
   136  		MessageType:    "WindowsBroadcast2",
   137  		RequiredLabels: []*fspb.Label{{ServiceName: "system", Label: "Windows"}},
   138  		Data: &anypb.Any{
   139  			TypeUrl: "message proto name 2",
   140  			Value:   []byte("message data 2"),
   141  		}}
   142  	if err := ds.CreateBroadcast(ctx, br2, 8); err != nil {
   143  		t.Fatalf("Unable to CreateBroadcast(%v): %v", br2, err)
   144  	}
   145  
   146  	time.Sleep(300 * time.Millisecond)
   147  
   148  	if err := bm.Close(ctx); err != nil {
   149  		t.Error(err)
   150  	}
   151  }