Golang slugid.Nice函数代码示例

func TestTaskContextConcurrentLogging(t *testing.T) {
	t.Parallel()
	path := filepath.Join(os.TempDir(), slugid.Nice())
	context, control, err := NewTaskContext(path, TaskInfo{}, nil)
	nilOrPanic(err, "Failed to create context")

	wg := sync.WaitGroup{}
	wg.Add(5) // This could trigger errors with race condition detector
	go func() { context.Log("Hello World 2"); wg.Done() }()
	go func() { context.Log("Hello World 1"); wg.Done() }()
	go func() { context.Log("Hello World 3 - Cheese"); wg.Done() }()
	go func() { context.Log("Hello World 4"); wg.Done() }()
	go func() { context.Log("Hello World 5"); wg.Done() }()
	wg.Wait()
	err = control.CloseLog()
	nilOrPanic(err, "Failed to close log file")

	reader, err := context.NewLogReader()
	nilOrPanic(err, "Failed to open log file")
	data, err := ioutil.ReadAll(reader)
	nilOrPanic(err, "Failed to read log file")

	if !strings.Contains(string(data), "Cheese") {
		panic("Couldn't find 'Cheese' in the log")
	}
	nilOrPanic(reader.Close(), "Failed to close log file")
	err = context.logStream.Remove()
	nilOrPanic(err, "Failed to remove logStream")
}

// New starts a livelog OS process using the executable specified, and returns
// a *LiveLog. The *LiveLog provides access to the GetURL which can be used to
// tail the log by multiple consumers in parallel, together with an
// io.WriteCloser where the logs should be written to. It is envisanged that
// the io.WriteCloser is passed on to the executing process.
//
// sslCert and sslKey should be used to specify the file location of a
// suitable certificate and key on the local filesystem that can be used
// for hosting the livelog service over https. If either is an empty string
// the livelog will resort to running over http transport instead.
//
// Please note the GetURL is for the loopback interface - it is beyond the
// scope of this library to transform this localhost URL into a URL with a
// fully qualified hostname using package
// github.com/taskcluster/stateless-dns-go/hostname since this package can be
// used independently of the former one.
func New(liveLogExecutable, sslCert, sslKey string) (*LiveLog, error) {
	l := &LiveLog{
		secret:  slugid.Nice(),
		command: exec.Command(liveLogExecutable),
		sslCert: sslCert,
		sslKey:  sslKey,
	}
	l.command.Env = append(
		os.Environ(),
		"ACCESS_TOKEN="+l.secret,
		"SERVER_CRT_FILE="+l.sslCert,
		"SERVER_KEY_FILE="+l.sslKey,
	)
	err := l.command.Start()
	// TODO: we need to make sure that this livelog process we just started
	// doesn't just exit, which can happen if the port is already in use!!!
	// Note, this is really bad, since another livelog will use a different
	// secret. Also note we get a 0 exit code when process exits because
	// another process was listening on the port(s).
	if err != nil {
		return nil, err
	}
	l.setRequestURLs()
	err = l.connectInputStream()
	// Note we can't wait for GET port to be active before returning
	// since livelog will only serve from that port once some content is
	// sent - so no good to execute waitForPortToBeActive(60023) here...
	// We would need to fix this in livelog codebase not here...
	return l, err
}

func (a artifactTestCase) Test() {
	taskID := slugid.Nice()
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "Hello, client.")
	}))
	defer ts.Close()

	s3resp, _ := json.Marshal(queue.S3ArtifactResponse{
		PutURL: ts.URL,
	})
	resp := queue.PostArtifactResponse(s3resp)
	mockedQueue := &client.MockQueue{}
	for _, path := range a.Artifacts {
		mockedQueue.On(
			"CreateArtifact",
			taskID,
			"0",
			path,
			client.PostS3ArtifactRequest,
		).Return(&resp, nil)
	}

	a.Case.QueueMock = mockedQueue
	a.Case.TaskID = taskID
	a.Case.Test()
	mockedQueue.AssertExpectations(a.Case.TestStruct)
}

// CreateUser will create a new user, with the given homeFolder, set the user
// owner of the homeFolder, and assign the user membership of given groups.
func CreateUser(homeFolder string, groups []*Group) (*User, error) {
	// Prepare arguments
	args := formatArgs(map[string]string{
		"-d": homeFolder,   // Set home folder
		"-c": "task user",  // Comment
		"-s": defaultShell, // Set default shell
	})
	args = append(args, "-M") // Don't create home, ignoring any global settings
	args = append(args, "-U") // Create primary user-group with same name
	if len(groups) > 0 {
		gids := []string{}
		for _, g := range groups {
			gids = append(gids, strconv.Itoa(g.gid))
		}
		args = append(args, "-G", strings.Join(gids, ","))
	}

	// Generate a random username
	name := slugid.Nice()
	args = append(args, name)

	// Run useradd command
	_, err := exec.Command(systemUserAdd, args...).Output()
	if err != nil {
		if e, ok := err.(*exec.ExitError); ok {
			return nil, fmt.Errorf(
				"Failed to create user with useradd, stderr: '%s'", string(e.Stderr),
			)
		}
		return nil, fmt.Errorf("Failed to run useradd, error: %s", err)
	}

	// Lookup user to get the uid
	u, err := user.Lookup(name)
	if err != nil {
		panic(fmt.Sprintf(
			"Failed to lookup newly created user: '%s', error: %s",
			name, err,
		))
	}

	// Parse uid/gid
	uid, err := strconv.ParseUint(u.Uid, 10, 32)
	if err != nil {
		panic(fmt.Sprintf("user.Uid should be an integer on POSIX systems"))
	}
	gid, err := strconv.ParseUint(u.Gid, 10, 32)
	if err != nil {
		panic(fmt.Sprintf("user.Gid should be an integer on POSIX systems"))
	}
	debug("Created user with uid: %d, gid: %d, name: %s", uid, gid, name)

	// Set user as owner of home folder
	err = os.Chown(homeFolder, int(uid), int(gid))
	if err != nil {
		return nil, fmt.Errorf("Failed to chown homeFolder, error: %s", err)
	}

	return &User{uint32(uid), uint32(gid), name, homeFolder}, nil
}

func TestNon200HasErrorBody(t *testing.T) {
	test := func(t *testing.T, creds *tcclient.Credentials) *httptest.ResponseRecorder {

		// Test setup
		routes := Routes{
			ConnectionData: tcclient.ConnectionData{
				Authenticate: true,
				Credentials:  creds,
			},
		}
		taskId := slugid.Nice()

		req, err := http.NewRequest(
			"POST",
			"http://localhost:60024/queue/v1/task/"+taskId+"/define",
			bytes.NewBufferString(
				`{"comment": "Valid json so that we hit endpoint, but should not result in http 200"}`,
			),
		)
		if err != nil {
			log.Fatal(err)
		}
		res := httptest.NewRecorder()

		// Function to test
		routes.ServeHTTP(res, req)

		// Validate results
		return res

	}
	testWithPermCreds(t, test, 400)
	testWithTempCreds(t, test, 400)
}

// Instance will return an Instance of the image with imageID. If no such
// image exists in the cache, download() will be called to download it to a
// temporary filename.
//
// This method will insert the downloaded image into the cache, and ensures that
// if won't be downloaded twice, if another invocation already is downloading
// an image with the same imageID.
//
// It is the responsibility of the caller to make sure that imageID is a string
// that uniquely identifies the image. Sane patterns includes "url:<url>", or
// "taskId:<taskId>/<runId>/<artifact>". It also the callers responsibility to
// enforce any sort of access control.
func (m *Manager) Instance(imageID string, download Downloader) (*Instance, error) {
	m.m.Lock()

	// Get image from cache and insert it if not present
	img := m.images[imageID]
	if img == nil {
		imageDone := make(chan struct{})
		img = &image{
			imageID: imageID,
			folder:  filepath.Join(m.imageFolder, slugid.Nice()),
			done:    imageDone,
			manager: m,
		}
		m.images[imageID] = img
		// Start loading the image
		go img.loadImage(download, imageDone)
	}

	// Acqure the image, so we can release lock without risking the image gets
	// garbage collected.
	img.Acquire()
	m.m.Unlock() // Release lock we don't need it anymore

	// Wait for image to be done, then either return the error, or return an
	// instance of the image.
	<-img.done
	if img.err != nil {
		img.Release()
		return nil, img.err
	}
	return img.instance()
}

func newMockDisplay() io.ReadWriteCloser {
	// Create listener for randomly generated addr
	addr := slugid.Nice()
	l, err := mocknet.Listen(addr)
	if err != nil {
		// This shouldn't be possible
		panic(fmt.Sprintf("mocknet.Listen failed using random addr, error: %s", err))
	}

	// Create and start server
	s := rfb.NewServer(mockDisplayWidth, mockDisplayHeight)
	go s.Serve(l)

	// Dial up to server
	conn, err := mocknet.Dial(addr)
	if err != nil {
		// This shouldn't happen either
		panic(fmt.Sprintf("mocknet.Dial failed, error: %s", err))
	}

	// Handle display when we get a connection from the server
	go handleDisplay(<-s.Conns) // This works because Conns has a size 16

	// Stop listener, we'll create one for each mock display connection
	l.Close()

	return conn
}

func ensureEnvironment(t *testing.T) (*runtime.Environment, engines.Engine, plugins.Plugin) {
	tempPath := filepath.Join(os.TempDir(), slugid.Nice())
	tempStorage, err := runtime.NewTemporaryStorage(tempPath)
	if err != nil {
		t.Fatal(err)
	}

	environment := &runtime.Environment{
		TemporaryStorage: tempStorage,
	}
	engineProvider := engines.Engines()["mock"]
	engine, err := engineProvider.NewEngine(engines.EngineOptions{
		Environment: environment,
		Log:         logger.WithField("engine", "mock"),
	})
	if err != nil {
		t.Fatal(err.Error())
	}

	pluginOptions := plugins.PluginOptions{
		Environment: environment,
		Engine:      engine,
		Log:         logger.WithField("component", "Plugin Manager"),
	}

	pm, err := plugins.Plugins()["success"].NewPlugin(pluginOptions)
	if err != nil {
		t.Fatalf("Error creating task manager. Could not create plugin manager. %s", err)
	}

	return environment, engine, pm
}

// NewUserNetwork returns a Network implementation using the QEMU user-space
// network stack. This doesn't provide the same level of isolation, but the
// meta-data service should be sufficiently isolated.
func NewUserNetwork(socketFolder string) (*UserNetwork, error) {
	n := &UserNetwork{
		socketFile: filepath.Join(socketFolder, "meta-"+slugid.Nice()+".sock"),
	}
	n.server = &graceful.Server{
		Timeout: 35 * time.Second,
		Server: &http.Server{
			Addr:    metaDataIP + ":80",
			Handler: http.HandlerFunc(n.dispatchRequest),
		},
		NoSignalHandling: true,
	}

	// Start listening (we handle listener error as a special thing)
	listener, err := net.ListenUnix("unix", &net.UnixAddr{
		Name: n.socketFile,
		Net:  "unix",
	})
	if err != nil {
		return nil, fmt.Errorf("Failed to listen on %s error: %s", n.socketFile, err)
	}

	// Start serving
	serverDone := make(chan struct{})
	n.serverDone = serverDone
	go func(n *UserNetwork, done chan<- struct{}) {
		err := n.server.Serve(listener)
		close(done)
		if err != nil {
			panic(fmt.Sprint("Fatal: meta-data service listener failed, error: ", err))
		}
	}(n, serverDone)

	return n, nil
}

func taskWithPayload(payload string) TaskRun {
	return TaskRun{
		TaskID: slugid.Nice(),
		Definition: queue.TaskDefinitionResponse{
			Payload: json.RawMessage(payload),
		},
		logWriter: &bytes.Buffer{},
	}
}

// Downloads ArtifactContent to a file inside the downloads directory specified
// in the global config file. The filename is a random slugid, and the
// absolute path of the file is returned.
func (ac *ArtifactContent) Download() (string, error) {
	basename := slugid.Nice()
	file := filepath.Join(config.DownloadsDir, basename)
	signedURL, err := Queue.GetLatestArtifact_SignedURL(ac.TaskID, ac.Artifact, time.Minute*30)
	if err != nil {
		return "", err
	}
	return file, downloadURLToFile(signedURL.String(), file)
}

func TestBuildImage(t *testing.T) {
	// Setup logging
	logger, _ := runtime.CreateLogger("info")
	log := logger.WithField("component", "qemu-build")

	inputImageFile, err := filepath.Abs("../../engines/qemu/test-image/tinycore-setup.tar.zst")
	if err != nil {
		panic(err)
	}
	outputFile := filepath.Join(os.TempDir(), slugid.Nice())
	defer os.Remove(outputFile)
	novnc := true
	cdrom := ""

	// Create ISO file to play with
	datadir := filepath.Join(os.TempDir(), slugid.Nice())
	defer os.RemoveAll(datadir)
	err = os.Mkdir(datadir, 0700)
	if err != nil {
		panic(err)
	}
	err = ioutil.WriteFile(filepath.Join(datadir, "setup.sh"),
		[]byte("#!/bin/sh\necho 'started';\nsudo poweroff;\n"), 0755)
	if err != nil {
		panic(err)
	}
	isofile := filepath.Join(os.TempDir(), slugid.Nice())
	defer os.Remove(isofile)
	err = exec.Command("genisoimage", "-vJrV", "DATA_VOLUME", "-input-charset", "utf-8", "-o", isofile, datadir).Run()
	if err != nil {
		panic(err)
	}

	err = buildImage(
		log, inputImageFile, outputFile,
		true, novnc, isofile, cdrom, 1,
	)
	if err != nil {
		panic(err)
	}
}

// instance returns a new instance of the image for use in a virtual machine.
// You must have called image.Acquire() first to prevent garbage collection.
func (img *image) instance() (*Instance, error) {
	// Create a copy of layer.qcow2
	diskFile := filepath.Join(img.folder, slugid.Nice()+".qcow2")
	err := copyFile(filepath.Join(img.folder, "layer.qcow2"), diskFile)
	if err != nil {
		return nil, fmt.Errorf("Failed to make copy of layer.qcow2, error: %s", err)
	}

	return &Instance{
		image:    img,
		diskFile: diskFile,
	}, nil
}

func (img *image) loadImage(download Downloader, done chan<- struct{}) {
	imageFile := filepath.Join(img.manager.imageFolder, slugid.Nice()+".tar.zst")

	// Create image folder
	err := os.Mkdir(img.folder, 0777)
	if err != nil {
		goto cleanup
	}

	// Download image to tempory file
	err = download(imageFile)
	if err != nil {
		goto cleanup
	}

	// Extract image and validate image
	img.machine, err = extractImage(imageFile, img.folder)
	if err != nil {
		goto cleanup
	}

	// Clean up if there is any error
cleanup:
	// Delete the image file
	e := os.RemoveAll(imageFile)
	if e != nil {
		eventID := img.manager.sentry.CaptureError(e, nil) // TODO: Severity level warning
		img.manager.log.Warning("Failed to delete image file, err: ", e, " sentry eventId: ", eventID)
	}

	// If there was an err, set ima.err and remove it from cache
	if err != nil {
		img.err = err
		// We should always remove a failed attempt from the cache
		img.manager.m.Lock()
		delete(img.manager.images, img.imageID)
		img.manager.m.Unlock()

		// Delete the image folder
		e := os.RemoveAll(img.folder)
		if e != nil {
			eventID := img.manager.sentry.CaptureError(e, nil) // TODO: Severity level warning
			img.manager.log.Warning("Failed to delete image folder, err: ", e, " sentry eventId: ", eventID)
		}
	} else {
		img.manager.gc.Register(img)
	}
	close(done)
}

// handlePoll handles GET /engine/v1/poll
func (s *MetaService) handlePoll(w http.ResponseWriter, r *http.Request) {
	if !forceMethod(w, r, http.MethodGet) {
		return
	}

	debug("GET /engine/v1/poll")
	select {
	case <-s.haltPolling:
		reply(w, http.StatusOK, Action{
			ID:   slugid.Nice(),
			Type: "none",
		})
	case <-time.After(PollTimeout):
		reply(w, http.StatusOK, Action{
			ID:   slugid.Nice(),
			Type: "none",
		})
	case action := <-s.actionOut:
		debug(" -> Sending action with id=%s", action.ID)
		if reply(w, http.StatusOK, action) != nil {
			debug("Failed to send action id=%s", action.ID)

			// Take the asyncRecord record out of the dictionary
			s.mPendingRecords.Lock()
			rec := s.pendingRecords[action.ID]
			delete(s.pendingRecords, action.ID)
			s.mPendingRecords.Unlock()

			// If nil, then the request is already being handled, and there is no need
			// to abort (presumably the action we received on the other side)
			if rec != nil {
				close(rec.Done)
			}
		}
	}
}