diff --git a/main.go b/main.go index 83af596..4d05ad2 100644 --- a/main.go +++ b/main.go @@ -1,58 +1,19 @@ package main import ( - "context" - "time" - "github.com/docker/docker/api/types" - "github.com/docker/docker/api/types/events" - "github.com/docker/docker/api/types/filters" - "github.com/docker/docker/client" "github.com/patrickmn/go-cache" "github.com/rs/zerolog/log" "github.com/shinebayar-g/ufw-docker-automated/logger" "github.com/shinebayar-g/ufw-docker-automated/ufwhandler" ) -func createClient() (*context.Context, *client.Client, error) { - ctx := context.Background() - client, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation()) - if err != nil { - return &ctx, client, err - } - _, err = client.Info(ctx) - return &ctx, client, err -} - -func streamEvents(ctx *context.Context, c *client.Client) (<-chan events.Message, <-chan error) { - filter := filters.NewArgs() - filter.Add("type", "container") - filter.Add("label", "UFW_MANAGED=TRUE") - return c.Events(*ctx, types.EventsOptions{Filters: filter}) -} - -func reconnect() (*context.Context, *client.Client) { - var ctx *context.Context - var client *client.Client - var err error - for { - time.Sleep(5 * time.Second) - log.Info().Msg("ufw-docker-automated: Trying to reconnect..") - ctx, client, err = createClient() - if err == nil { - break - } - } - log.Info().Msg("ufw-docker-automated: Reconnected to the Docker Engine.") - return ctx, client -} - func main() { logger.SetupLogger() - ctx, client, err := createClient() + ctx, client, err := ufwhandler.CreateClient() if err != nil { log.Error().Err(err).Msg("ufw-docker-automated: Client error.") - ctx, client = reconnect() + ctx, client = ufwhandler.Reconnect() } else { log.Info().Msg("ufw-docker-automated: Connected to the Docker Engine.") } @@ -65,7 +26,7 @@ func main() { go ufwhandler.Cleanup(ctx, client) go ufwhandler.Sync(ctx, createChannel, client) - messages, errors := streamEvents(ctx, client) + messages, errors := ufwhandler.StreamEvents(ctx, client) for { select { case msg := <-messages: @@ -83,9 +44,9 @@ func main() { case err := <-errors: if err != nil { log.Error().Err(err).Msg("ufw-docker-automated: Event error.") - ctx, client = reconnect() + ctx, client = ufwhandler.Reconnect() go ufwhandler.Sync(ctx, createChannel, client) - messages, errors = streamEvents(ctx, client) + messages, errors = ufwhandler.StreamEvents(ctx, client) } } } diff --git a/ufwhandler/client.go b/ufwhandler/client.go new file mode 100644 index 0000000..e188863 --- /dev/null +++ b/ufwhandler/client.go @@ -0,0 +1,45 @@ +package ufwhandler + +import ( + "context" + "time" + + "github.com/docker/docker/api/types" + "github.com/docker/docker/api/types/events" + "github.com/docker/docker/api/types/filters" + "github.com/docker/docker/client" + "github.com/rs/zerolog/log" +) + +func CreateClient() (*context.Context, *client.Client, error) { + ctx := context.Background() + client, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation()) + if err != nil { + return &ctx, client, err + } + _, err = client.Info(ctx) + return &ctx, client, err +} + +func Reconnect() (*context.Context, *client.Client) { + var ctx *context.Context + var client *client.Client + var err error + for { + time.Sleep(5 * time.Second) + log.Info().Msg("ufw-docker-automated: Trying to reconnect..") + ctx, client, err = CreateClient() + if err == nil { + break + } + } + log.Info().Msg("ufw-docker-automated: Reconnected to the Docker Engine.") + return ctx, client +} + +func StreamEvents(ctx *context.Context, c *client.Client) (<-chan events.Message, <-chan error) { + filter := filters.NewArgs() + filter.Add("type", "container") + filter.Add("label", "UFW_MANAGED=TRUE") + return c.Events(*ctx, types.EventsOptions{Filters: filter}) +}