create.go

package main

import (
	"bytes"
	"context"
	"encoding/base64"
	"fmt"
	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/route53"
	"github.com/aws/aws-sdk-go-v2/service/route53/types"
	"github.com/spf13/cobra"
	"io"
	"os"
	"strings"
	"time"
)

// This, with some math equates to exactly 4000 characters generated by base64. This is the exact limit
// of aws route53 TXT record sizes
const maxBase64ReadSize = 2964
const maxTXTRRSize int = 3952 // It's really 4000, but for us and padding problems, it's 3998, but actual data we can store is ~3952

// createCmd represents the create command
var createCmd = &cobra.Command{
	Use:   "create",
	Short: "Create a \"CDN\" on route53",
	Long:  `Read the provided file, encode it, and store it in a definable pattern on route53 via TXT records`,
	PreRun: func(cmd *cobra.Command, args []string) {
		fileLocation, err := cmd.Flags().GetString("file")
		if err != nil {
			panic("Impossible that --file doesn't have a value")
		}

		if fileLocation == "-" {
			targetFile = os.Stdin
		} else {
			targetFile, err = os.Open(fileLocation)
			if err != nil {
				_, _ = fmt.Fprintf(os.Stderr, "Failed to open file %v: %v\n", fileLocation, err)
				os.Exit(1)
			}
		}
	},
	Run: func(cmd *cobra.Command, args []string) {
		createCDN()
	},
}

func createCDN() {
	client := getClient()

	targetZone := getHostedZone(client)
	if targetZone != nil {
		_, _ = fmt.Fprintf(os.Stderr, "Found target zone %v\n", *targetZone.Id)
	} else {
		_, _ = fmt.Fprintf(os.Stderr, "Failed to find hosted zone with name %v\n", specificHostedZoneDomain)
		os.Exit(1)
	}

	RRSChan := make(chan *types.ResourceRecordSet, 100)
	uploadDone := make(chan bool, 0)

	// Write to aws
	go func() {
		count := uint64(0)
		changeSet := &route53.ChangeResourceRecordSetsInput{
			ChangeBatch: &types.ChangeBatch{
				Changes: make([]types.Change, 0),
			},
			HostedZoneId: targetZone.Id,
		}

		for {
			newRRS, ok := <-RRSChan
			if ok || len(changeSet.ChangeBatch.Changes) > 0 {
				if newRRS != nil {
					count++
					newRRS.Name = aws.String(fmt.Sprintf("%s-%v.%s", specificSubdomain, count, specificHostedZoneDomain))
					newRRS.TTL = aws.Int64(60)

					changeSet.ChangeBatch.Changes = append(changeSet.ChangeBatch.Changes, types.Change{
						Action:            "CREATE",
						ResourceRecordSet: newRRS,
					})
				}

				if len(changeSet.ChangeBatch.Changes) >= 5 || !ok {
					_, err := client.ChangeResourceRecordSets(context.Background(), changeSet)
					if err != nil {
						_, _ = fmt.Fprintf(os.Stderr, "Failed to submit RRS change to AWS: %v\n", err)
						os.Exit(1)
					}

					changeSet.ChangeBatch.Changes = make([]types.Change, 0)
				}
			} else { // channel closed right?
				_, err := client.ChangeResourceRecordSets(context.Background(), &route53.ChangeResourceRecordSetsInput{
					ChangeBatch: &types.ChangeBatch{
						Changes: []types.Change{
							{
								Action: "CREATE",
								ResourceRecordSet: &types.ResourceRecordSet{
									Name: aws.String(fmt.Sprintf("%s-0.%s", specificSubdomain, specificHostedZoneDomain)),
									Type: "TXT",
									ResourceRecords: []types.ResourceRecord{
										{
											Value: aws.String(fmt.Sprintf("\"%v\"", count)),
										},
									},
									TTL: aws.Int64(60),
								},
							},
						},
						Comment: nil,
					},
					HostedZoneId: targetZone.Id,
				})

				if err != nil {
					_, _ = fmt.Fprintf(os.Stderr, "Failed to submit RRS change to AWS: %v\n", err)
					os.Exit(1)
				}

				uploadDone <- true
				break
			}
		}
	}()

	middleware := &bytes.Buffer{}
	b64Encoder := base64.NewEncoder(base64.StdEncoding, middleware)
	rrValueBuffer := &strings.Builder{}
	uploading := uint64(0)
	start := time.Now()
	for numWrittenToMiddleware, err := io.CopyN(b64Encoder, targetFile, maxBase64ReadSize); ; numWrittenToMiddleware, err = io.CopyN(b64Encoder, targetFile, maxBase64ReadSize) {
		exit := false
		if err != nil {
			exit = true // Save our error and createCDN one last time before exit
			err = nil
		}

		if numWrittenToMiddleware > maxBase64ReadSize {
			panic(fmt.Sprintf("We wrote too much and have broken our encoding segmentation: %v", numWrittenToMiddleware))
		}

		b64Encoder.Close() // Flush
		b64Encoder = base64.NewEncoder(base64.StdEncoding, middleware)

		// Maximum number of bytes we can have in a TXT Answer
		if middleware.Len() > maxTXTRRSize {
			fmt.Fprintf(os.Stderr, "middleware holds %v when wanted <= %v", middleware.Len(), maxTXTRRSize)
			os.Exit(1)
		}

		encodedData := middleware.Next(maxTXTRRSize)
		uploading += uint64(len(encodedData))

		for {
			if len(encodedData) <= 255 {
				rrValueBuffer.WriteString(fmt.Sprintf("\"%s\" ", encodedData))
				break
			} else {
				rrValueBuffer.WriteString(fmt.Sprintf("\"%s\" ", encodedData[:255]))
				encodedData = encodedData[255:]
			}
		}

		newRRS := &types.ResourceRecordSet{
			Type: "TXT",
			ResourceRecords: []types.ResourceRecord{
				{
					Value: aws.String(rrValueBuffer.String()),
				},
			},
		}

		if rrValueBuffer.Len() > 0 {
			RRSChan <- newRRS

			fmt.Fprintf(os.Stdout, "\r\033[2KUploading %v Bytes", uploading)

			duration := time.Now().Sub(start) / time.Second
			if duration >= 1 {
				rate := uploading / uint64(duration) / 1000
				fmt.Fprintf(os.Stdout, " (~%v KB/s)", rate)
			}
		}

		middleware.Reset()
		rrValueBuffer.Reset()

		if exit {
			break
		}
	}

	close(RRSChan)
	fmt.Println()
	<-uploadDone
}

func init() {
	awsCmd.AddCommand(createCmd)
}