-
Notifications
You must be signed in to change notification settings - Fork 11
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
7 changed files
with
270 additions
and
18 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,27 +1,80 @@ | ||
package plugininterface | ||
|
||
import ( | ||
"context" | ||
"net/rpc" | ||
|
||
log "github.com/sirupsen/logrus" | ||
|
||
"github.com/google/uuid" | ||
"github.com/trento-project/agent/pkg/factsengine/entities" | ||
) | ||
|
||
type GathererRPC struct{ client *rpc.Client } | ||
|
||
func (g *GathererRPC) Gather(factsRequest []entities.FactRequest) ([]entities.Fact, error) { | ||
func (g *GathererRPC) RequestGathering( | ||
ctx context.Context, | ||
factsRequest []entities.FactRequest, | ||
) ([]entities.Fact, error) { | ||
var resp []entities.Fact | ||
var err error | ||
|
||
requestID := uuid.New().String() | ||
args := GatheringArgs{ | ||
FactRequests: factsRequest, | ||
RequestID: requestID, | ||
} | ||
|
||
err := g.client.Call("Plugin.Gather", factsRequest, &resp) | ||
gathering := make(chan error) | ||
|
||
return resp, err | ||
go func() { | ||
gathering <- g.client.Call("Plugin.ServeGathering", args, &resp) | ||
}() | ||
|
||
select { | ||
case <-ctx.Done(): | ||
err = g.client.Call("Plugin.Cancel", requestID, &resp) | ||
return []entities.Fact{}, err | ||
case err = <-gathering: | ||
if err != nil { | ||
return nil, err | ||
} | ||
return resp, nil | ||
} | ||
} | ||
|
||
type GathererRPCServer struct { | ||
Impl Gatherer | ||
Impl Gatherer | ||
cancelMap map[string]context.CancelFunc | ||
} | ||
|
||
type GatheringArgs struct { | ||
FactRequests []entities.FactRequest | ||
RequestID string | ||
} | ||
|
||
func (s *GathererRPCServer) Gather(args []entities.FactRequest, resp *[]entities.Fact) error { | ||
func (s *GathererRPCServer) ServeGathering(args GatheringArgs, resp *[]entities.Fact) error { | ||
|
||
ctx, cancel := context.WithCancel(context.Background()) | ||
if s.cancelMap == nil { | ||
s.cancelMap = make(map[string]context.CancelFunc) | ||
} | ||
s.cancelMap[args.RequestID] = cancel | ||
defer delete(s.cancelMap, args.RequestID) | ||
|
||
var err error | ||
*resp, err = s.Impl.Gather(args) | ||
*resp, err = s.Impl.Gather(ctx, args.FactRequests) | ||
return err | ||
} | ||
|
||
func (s *GathererRPCServer) Cancel(requestID string, _ *[]entities.Fact) (_ error) { | ||
cancel, ok := s.cancelMap[requestID] | ||
if ok { | ||
cancel() | ||
delete(s.cancelMap, requestID) | ||
} else { | ||
log.Warnf("Cannot find cancel function for request %s", requestID) | ||
} | ||
|
||
return nil | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,68 @@ | ||
package main | ||
|
||
// go build -o /usr/etc/trento/sleep ./plugin_examples/sleep/sleep.go | ||
|
||
import ( | ||
"context" | ||
"fmt" | ||
"os/exec" | ||
"sync" | ||
|
||
"github.com/hashicorp/go-plugin" | ||
log "github.com/sirupsen/logrus" | ||
"github.com/trento-project/agent/pkg/factsengine/entities" | ||
"github.com/trento-project/agent/pkg/factsengine/plugininterface" | ||
) | ||
|
||
type sleepGatherer struct { | ||
} | ||
|
||
func (s sleepGatherer) Gather(ctx context.Context, factsRequests []entities.FactRequest) ([]entities.Fact, error) { | ||
facts := []entities.Fact{} | ||
|
||
log.Infof("Starting sleep plugin facts gathering process") | ||
|
||
wg := sync.WaitGroup{} | ||
|
||
for _, factReq := range factsRequests { | ||
log.Infof("Sleeping for %s", factReq.Argument) | ||
fact := entities.NewFactGatheredWithRequest(factReq, &entities.FactValueString{Value: fmt.Sprint(factReq.Argument)}) | ||
facts = append(facts, fact) | ||
|
||
time := fmt.Sprint(factReq.Argument) | ||
wg.Add(1) | ||
go func(time string) { | ||
defer wg.Done() | ||
cmd := exec.CommandContext(ctx, "sleep", time) | ||
err := cmd.Run() | ||
if err != nil { | ||
log.Errorf("Error running sleep command: %s", err) | ||
} | ||
}(time) | ||
|
||
} | ||
|
||
wg.Wait() | ||
|
||
log.Infof("Requested sleep plugin facts gathered") | ||
return facts, nil | ||
} | ||
|
||
func main() { | ||
d := &sleepGatherer{} | ||
|
||
handshakeConfig := plugin.HandshakeConfig{ | ||
ProtocolVersion: 1, | ||
MagicCookieKey: "TRENTO_PLUGIN", | ||
MagicCookieValue: "gatherer", | ||
} | ||
|
||
var pluginMap = map[string]plugin.Plugin{ | ||
"gatherer": &plugininterface.GathererPlugin{Impl: d}, | ||
} | ||
|
||
plugin.Serve(&plugin.ServeConfig{ // nolint | ||
HandshakeConfig: handshakeConfig, | ||
Plugins: pluginMap, | ||
}) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters