ipam.go

/*
 * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package ipam

import (
	"context"
	"errors"
	"fmt"
	"net/netip"
	"strings"

	cipam "github.com/nvidia/bare-metal-manager-rest/ipam"
	"github.com/uptrace/bun"
	"go4.org/netipx"

	cdb "github.com/nvidia/bare-metal-manager-rest/db/pkg/db"
	cdbm "github.com/nvidia/bare-metal-manager-rest/db/pkg/db/model"
)

var (
	// ErrPrefixDoesNotExistForIPBlock is the error returned when ipam does not have the entry for the IPBlock
	ErrPrefixDoesNotExistForIPBlock = errors.New("prefix does not exist for IPBlock in ipam db")
	// ErrNilIPBlock is the error when a nil IPBlock was passed
	ErrNilIPBlock = errors.New("ipblock parameter is nil")
)

// ~~~~~ IPAM Utilities ~~~~~ //

// NewIpamStorage will return a bun ipam storage interface
func NewIpamStorage(db *bun.DB, tx *bun.Tx) cipam.Storage {
	return cipam.NewBunStorage(db, tx)
}

// GetFirstIPFromCidr will parse a cidr, and returns the first IP address in that cidr
// this is used as the gateway IP
func GetFirstIPFromCidr(cidr string) (string, error) {
	ipPref, err := netip.ParsePrefix(cidr)
	if err != nil {
		return "", err
	}
	return ipPref.Addr().Next().String(), nil
}

// ParseCidrIntoPrefixAndBlockSize will parse a cidr into the masked prefix, and blocksize
func ParseCidrIntoPrefixAndBlockSize(cidr string) (string, int, error) {
	ipPref, err := netip.ParsePrefix(cidr)
	if err != nil {
		return "", 0, err
	}
	return ipPref.Masked().Addr().String(), int(ipPref.Bits()), nil
}

// GetIpamNamespaceForIPBlock will return the namespace string for the IPBlock
// namespace is currently: routingtype/infrastructureProviderID/siteID
func GetIpamNamespaceForIPBlock(ctx context.Context, routingType string, infrastructureProviderID, siteID string) string {
	return fmt.Sprintf("%s/%s/%s", routingType, infrastructureProviderID, siteID)
}

// GetCidrForIPBlock will return the cidr given the prefix, and block size
func GetCidrForIPBlock(ctx context.Context, prefix string, blockSize int) string {
	return fmt.Sprintf("%s/%d", prefix, blockSize)
}

// CreateIpamEntryForIPBlock will create an ipam entry in the ipam DB for the IPBlock
// will error if there is a prefix clash in that same namespace
func CreateIpamEntryForIPBlock(ctx context.Context, ipamDB cipam.Storage, prefix string, blockSize int, routingType string, infrastructureProviderID, siteID string) (*cipam.Prefix, error) {
	ipamer := cipam.NewWithStorage(ipamDB)
	namespace := GetIpamNamespaceForIPBlock(ctx, routingType, infrastructureProviderID, siteID)
	ipamer.SetNamespace(namespace)
	cidr := GetCidrForIPBlock(ctx, prefix, blockSize)
	ipamPrefix, err := ipamer.NewPrefix(ctx, cidr)
	if err != nil {
		return nil, err
	}
	return ipamPrefix, err
}

// DeleteIpamEntryForIPBlock will delete the ipam entry in the ipam DB for the IPBlock
// will not error if the deleted cidr is not existing (idempotent)
func DeleteIpamEntryForIPBlock(ctx context.Context, ipamDB cipam.Storage, prefix string, blockSize int, routingType string, infrastructureProviderID, siteID string) error {
	ipamer := cipam.NewWithStorage(ipamDB)
	namespace := GetIpamNamespaceForIPBlock(ctx, routingType, infrastructureProviderID, siteID)
	ipamer.SetNamespace(namespace)
	cidr := GetCidrForIPBlock(ctx, prefix, blockSize)
	_, err := ipamer.DeletePrefix(ctx, cidr)
	if err != nil {
		// TODO - encapsulate errors with types in the ipam package
		if strings.HasPrefix(err.Error(), cipam.ErrNotFound.Error()) {
			// if not found, dont consider it an error
			return nil
		}
		return err
	}
	return nil
}

// GetIpamUsageForIPBlock will get an ipam usage for the IPBlock
func GetIpamUsageForIPBlock(ctx context.Context, ipamDB cipam.Storage, ipBlock *cdbm.IPBlock) (*cipam.Usage, error) {

	if ipBlock == nil {
		return nil, ErrNilIPBlock
	}

	ipamer := cipam.NewWithStorage(ipamDB)
	namespace := GetIpamNamespaceForIPBlock(ctx, ipBlock.RoutingType, ipBlock.InfrastructureProviderID.String(), ipBlock.SiteID.String())
	ipamer.SetNamespace(namespace)
	cidr := GetCidrForIPBlock(ctx, ipBlock.Prefix, ipBlock.PrefixLength)
	ipamPrefix := ipamer.PrefixFrom(ctx, cidr)

	if ipamPrefix == nil {
		return nil, errors.New(fmt.Sprintf("did not find prefix for IPBlock: %s", ipBlock.ID.String()))
	}

	// Handle full grant scenario
	if ipBlock.FullGrant {
		return &cipam.Usage{
			AvailableIPs:              0,
			AcquiredIPs:               0,
			AvailableSmallestPrefixes: 0,
			AvailablePrefixes:         nil,
			AcquiredPrefixes:          1,
		}, nil
	} else {
		return &cipam.Usage{
			AvailableIPs:              ipamPrefix.Usage().AvailableIPs,
			AcquiredIPs:               ipamPrefix.Usage().AcquiredIPs,
			AvailableSmallestPrefixes: ipamPrefix.Usage().AvailableSmallestPrefixes,
			AvailablePrefixes:         ipamPrefix.Usage().AvailablePrefixes,
			AcquiredPrefixes:          ipamPrefix.Usage().AcquiredPrefixes,
		}, nil
	}
}

// CreateChildIpamEntryForIPBlock will create an child ipam entry in the ipam DB for the given parent IP Block, with a given child block size
// Note: FullGrant is a special case when the childBlockSize matches the parentIPBlock, and the parentIPBlock has no
// child prefixes, then, the parentIPBlock is updated as a full grant in db, and its prefix is
// returned (without any updates to the ipam DB)
func CreateChildIpamEntryForIPBlock(ctx context.Context, tx *cdb.Tx, dbSession *cdb.Session, ipamDB cipam.Storage, parentIPBlock *cdbm.IPBlock, childBlockSize int) (*cipam.Prefix, error) {
	if parentIPBlock == nil {
		return nil, ErrNilIPBlock
	}
	// FullGrant of the parent IPBlock is also handled here to keep it localized so,
	// we can reason better wrt correctness.
	// TODO: look into implementing full grant in cloud-ipam library.
	if parentIPBlock.FullGrant {
		return nil, errors.New(fmt.Sprintf("parent IPBlock : %s already has a full-grant", parentIPBlock.ID.String()))
	}
	ipamer := cipam.NewWithStorage(ipamDB)
	namespace := GetIpamNamespaceForIPBlock(ctx, parentIPBlock.RoutingType, parentIPBlock.InfrastructureProviderID.String(), parentIPBlock.SiteID.String())
	ipamer.SetNamespace(namespace)
	parentCidr := GetCidrForIPBlock(ctx, parentIPBlock.Prefix, parentIPBlock.PrefixLength)

	if childBlockSize == parentIPBlock.PrefixLength {
		parentPrefix := ipamer.PrefixFrom(ctx, parentCidr)
		if parentPrefix == nil {
			return nil, errors.New(fmt.Sprintf("did not find prefix for parentIPBlock: %s", parentIPBlock.ID.String()))
		}
		parentUsage := parentPrefix.Usage()
		if parentUsage.AcquiredPrefixes > 0 {
			return nil, errors.New("parent IPBlock has allocated prefixes, cannot do a full-grant")
		}
		// mark the parentIPBlock as a full grant, and return the prefix corresponding to it
		ipbDAO := cdbm.NewIPBlockDAO(dbSession)
		_, err := ipbDAO.Update(
			ctx,
			tx,
			cdbm.IPBlockUpdateInput{
				IPBlockID: parentIPBlock.ID,
				FullGrant: cdb.GetBoolPtr(true),
			},
		)
		if err != nil {
			return nil, errors.New("unable to update parent IPBlock full-grant field")
		}
		parentIPBlock.FullGrant = true
		return parentPrefix, nil
	}
	childPrefix, err := ipamer.AcquireChildPrefix(ctx, parentCidr, uint8(childBlockSize))
	if err != nil {
		return nil, err
	}
	return childPrefix, err
}

// DeleteChildIpamEntryFromCidr will delete a child ipam entry in the ipam DB
// given the parent IPBlock, and child cidr
// Note: FullGrant is a special case when the parentIPBlock has a full grant, and the child
// is being deleted, then the parent IPBlock's full grant is cleared in db
func DeleteChildIpamEntryFromCidr(ctx context.Context, tx *cdb.Tx, dbSession *cdb.Session, ipamDB cipam.Storage, parentIPBlock *cdbm.IPBlock, childCidr string) error {
	if parentIPBlock == nil {
		return ErrNilIPBlock
	}
	// if parentIPBlock is a full grant, clear the parent's full grant
	// and return
	if parentIPBlock.FullGrant {
		parentCidr := GetCidrForIPBlock(ctx, parentIPBlock.Prefix, parentIPBlock.PrefixLength)
		// this is a consistency check
		if parentCidr != childCidr {
			// this should never happen, ie, in a full grant, parent cidr and child cidr should match
			return errors.New(fmt.Sprintf("parent IPBlock has full-grant, but childCidr: %s does not match parentCidr: %s", childCidr, parentCidr))
		}
		ipbDAO := cdbm.NewIPBlockDAO(dbSession)
		_, err := ipbDAO.Update(
			ctx,
			tx,
			cdbm.IPBlockUpdateInput{
				IPBlockID: parentIPBlock.ID,
				FullGrant: cdb.GetBoolPtr(false),
			},
		)

		if err != nil {
			return errors.New(fmt.Sprintf("unable to update IPBlock's full-grant, ipblock id: %s ", parentIPBlock.ID.String()))
		}
		parentIPBlock.FullGrant = false
		return nil
	}
	ipamer := cipam.NewWithStorage(ipamDB)
	namespace := GetIpamNamespaceForIPBlock(ctx, parentIPBlock.RoutingType, parentIPBlock.InfrastructureProviderID.String(), parentIPBlock.SiteID.String())
	ipamer.SetNamespace(namespace)
	prefix := ipamer.PrefixFrom(ctx, childCidr)
	if prefix == nil {
		return ErrPrefixDoesNotExistForIPBlock
	}
	err := ipamer.ReleaseChildPrefix(ctx, prefix)
	if err != nil {
		return err
	}
	return nil
}

// ValidateIPAddresses checks that none of the given IP addresses are reserved
// network addresses (network address, broadcast address, or gateway) for the
// specified CIDR prefix. It returns the list of valid IPs and the list of
// rejected IPs with reasons.
func ValidateIPAddresses(ipAddresses []string, cidr string, gateway string) (valid []string, rejected []string) {
	if len(ipAddresses) == 0 {
		return ipAddresses, nil
	}

	prefix, err := netip.ParsePrefix(cidr)
	if err != nil {
		// If we can't parse the prefix, we can't validate — pass through
		return ipAddresses, nil
	}

	iprange := netipx.RangeOfPrefix(prefix)
	networkAddr := iprange.From()
	broadcastAddr := iprange.To()

	var gatewayAddr netip.Addr
	if gateway != "" {
		gatewayAddr, _ = netip.ParseAddr(gateway)
	}

	for _, ipStr := range ipAddresses {
		addr, err := netip.ParseAddr(ipStr)
		if err != nil {
			rejected = append(rejected, fmt.Sprintf("%s (invalid IP)", ipStr))
			continue
		}

		if !prefix.Contains(addr) {
			rejected = append(rejected, fmt.Sprintf("%s (not in prefix %s)", ipStr, cidr))
			continue
		}

		if addr == networkAddr {
			rejected = append(rejected, fmt.Sprintf("%s (network address of %s)", ipStr, cidr))
			continue
		}

		if prefix.Addr().Is4() && addr == broadcastAddr {
			rejected = append(rejected, fmt.Sprintf("%s (broadcast address of %s)", ipStr, cidr))
			continue
		}

		if gatewayAddr.IsValid() && addr == gatewayAddr {
			rejected = append(rejected, fmt.Sprintf("%s (gateway address of %s)", ipStr, cidr))
			continue
		}

		valid = append(valid, ipStr)
	}

	return valid, rejected
}

// GetInterfaceCIDRAndGateway returns the CIDR and gateway for an interface
// based on its associated Subnet or VpcPrefix.
func GetInterfaceCIDRAndGateway(ifc *cdbm.Interface) (cidr string, gateway string) {
	if ifc.Subnet != nil {
		if ifc.Subnet.IPv4Prefix != nil {
			cidr = fmt.Sprintf("%s/%d", *ifc.Subnet.IPv4Prefix, ifc.Subnet.PrefixLength)
			if ifc.Subnet.IPv4Gateway != nil {
				gateway = *ifc.Subnet.IPv4Gateway
			}
		} else if ifc.Subnet.IPv6Prefix != nil {
			cidr = fmt.Sprintf("%s/%d", *ifc.Subnet.IPv6Prefix, ifc.Subnet.PrefixLength)
			if ifc.Subnet.IPv6Gateway != nil {
				gateway = *ifc.Subnet.IPv6Gateway
			}
		}
	} else if ifc.VpcPrefix != nil {
		// VpcPrefix.Prefix may already contain the prefix length (e.g. "192.172.0.0/24")
		// from the IPAM library's Cidr format. Parse it to extract just the address.
		if p, err := netip.ParsePrefix(ifc.VpcPrefix.Prefix); err == nil {
			cidr = p.String()
		} else {
			cidr = fmt.Sprintf("%s/%d", ifc.VpcPrefix.Prefix, ifc.VpcPrefix.PrefixLength)
		}
	}
	return cidr, gateway
}