Link Server

package main

import (
	"context"
	"crypto/rand"
	"fmt"
	"log"
	"math/big"
	"net"
	"sync"
	"sync/atomic"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/reflection"
	"google.golang.org/grpc/status"

	"shortener/pb"
)

type linkServer struct {
	pb.UnimplementedLinkServiceServer
	mu     sync.RWMutex
	links  map[string]*pb.Link
	nextID atomic.Int64
}

func newLinkServer() *linkServer {
	return &linkServer{links: make(map[string]*pb.Link)}
}

func (s *linkServer) CreateLink(ctx context.Context, req *pb.CreateLinkRequest) (*pb.CreateLinkResponse, error) {
	if req.GetUrl() == "" {
		return nil, status.Error(codes.InvalidArgument, "url is required")
	}
	id := s.nextID.Add(1)
	code := shortCode()
	link := &pb.Link{Id: id, Url: req.GetUrl(), ShortCode: code}
	s.mu.Lock()
	s.links[code] = link
	s.mu.Unlock()
	return &pb.CreateLinkResponse{Link: link}, nil
}

func (s *linkServer) GetLink(ctx context.Context, req *pb.GetLinkRequest) (*pb.GetLinkResponse, error) {
	s.mu.RLock()
	link, ok := s.links[req.GetShortCode()]
	s.mu.RUnlock()
	if !ok {
		return nil, status.Errorf(codes.NotFound, "link not found: %s", req.GetShortCode())
	}
	return &pb.GetLinkResponse{Link: link}, nil
}

func (s *linkServer) DeleteLink(ctx context.Context, req *pb.DeleteLinkRequest) (*pb.DeleteLinkResponse, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.links[req.GetShortCode()]; !ok {
		return nil, status.Errorf(codes.NotFound, "link not found: %s", req.GetShortCode())
	}
	delete(s.links, req.GetShortCode())
	return &pb.DeleteLinkResponse{}, nil
}

func shortCode() string {
	const chars = "abcdefghijklmnopqrstuvwxyz0123456789"
	b := make([]byte, 6)
	for i := range b {
		n, _ := rand.Int(rand.Reader, big.NewInt(int64(len(chars))))
		b[i] = chars[n.Int64()]
	}
	return string(b)
}

func main() {
	lis, err := net.Listen("tcp", ":50051")
	if err != nil {
		log.Fatal(err)
	}
	srv := grpc.NewServer()
	pb.RegisterLinkServiceServer(srv, newLinkServer())
	reflection.Register(srv)
	fmt.Println("gRPC server on :50051")
	log.Fatal(srv.Serve(lis))
}