2 files changed, 29 insertions, 26 deletions

diff --git a/src/netlink_delinearize.c b/src/netlink_delinearize.c
index a8f700aa..d75df215 100644
--- a/src/netlink_delinearize.c
+++ b/src/netlink_delinearize.c
@@ -408,14 +408,13 @@ static void netlink_parse_nat(struct netlink_parse_ctx *ctx,
 
 	reg2 = nfnl_nft_nat_get_sreg_proto_max(nle);
 	if (reg2 && reg2 != reg1) {
-		proto = netlink_get_register(ctx, loc, reg1);
+		proto = netlink_get_register(ctx, loc, reg2);
 		if (proto == NULL)
 			return netlink_error(ctx, loc,
 					     "NAT statement has no proto "
 					     "expression");
 
 		expr_set_type(proto, &inet_service_type, BYTEORDER_BIG_ENDIAN);
-		stmt->nat.proto = proto;
 		if (stmt->nat.proto != NULL)
 			proto = range_expr_alloc(loc, stmt->nat.proto, proto);
 		stmt->nat.proto = proto;
@@ -629,6 +628,10 @@ static void expr_postprocess(struct rule_pp_ctx *ctx,
 			expr->len = len;
 		}
 		break;
+	case EXPR_RANGE:
+		expr_postprocess(ctx, stmt, &expr->left);
+		expr_postprocess(ctx, stmt, &expr->right);
+		break;
 	case EXPR_SET_REF:
 	case EXPR_EXTHDR:
 	case EXPR_META:
diff --git a/src/netlink_linearize.c b/src/netlink_linearize.c
index f3ff7d14..a3841fb9 100644
--- a/src/netlink_linearize.c
+++ b/src/netlink_linearize.c
@@ -500,53 +500,53 @@ static void netlink_gen_nat_stmt(struct netlink_linearize_ctx *ctx,
 	struct nfnl_nft_expr *nle;
 	enum nft_registers amin_reg, amax_reg;
 	enum nft_registers pmin_reg, pmax_reg;
+	int registers = 0;
 
 	nle = alloc_nft_expr(nfnl_nft_nat_init);
 	nfnl_nft_nat_set_type(nle, stmt->nat.type);
 
 	if (stmt->nat.addr) {
-		switch (stmt->nat.addr->ops->type) {
-		default:
-			amin_reg = amax_reg = get_register(ctx);
-			netlink_gen_expr(ctx, stmt->nat.addr, amin_reg);
-			nfnl_nft_nat_set_sreg_addr_min(nle, amin_reg);
-			release_register(ctx);
-			break;
-		case EXPR_RANGE:
-			amin_reg = get_register(ctx);
+		amin_reg = get_register(ctx);
+		registers++;
+
+		if (stmt->nat.addr->ops->type == EXPR_RANGE) {
 			amax_reg = get_register(ctx);
+			registers++;
+
 			netlink_gen_expr(ctx, stmt->nat.addr->left, amin_reg);
 			netlink_gen_expr(ctx, stmt->nat.addr->right, amax_reg);
 			nfnl_nft_nat_set_sreg_addr_min(nle, amin_reg);
 			nfnl_nft_nat_set_sreg_addr_max(nle, amax_reg);
-			release_register(ctx);
-			release_register(ctx);
-			break;
+		} else {
+			netlink_gen_expr(ctx, stmt->nat.addr, amin_reg);
+			nfnl_nft_nat_set_sreg_addr_min(nle, amin_reg);
 		}
 
 	}
 
 	if (stmt->nat.proto) {
-		switch (stmt->nat.proto->ops->type) {
-		default:
-			pmin_reg = pmax_reg = get_register(ctx);
-			netlink_gen_expr(ctx, stmt->nat.proto, pmin_reg);
-			nfnl_nft_nat_set_sreg_proto_min(nle, pmin_reg);
-			release_register(ctx);
-			break;
-		case EXPR_RANGE:
-			pmin_reg = get_register(ctx);
+		pmin_reg = get_register(ctx);
+		registers++;
+
+		if (stmt->nat.proto->ops->type == EXPR_RANGE) {
 			pmax_reg = get_register(ctx);
+			registers++;
+
 			netlink_gen_expr(ctx, stmt->nat.proto->left, pmin_reg);
 			netlink_gen_expr(ctx, stmt->nat.proto->right, pmax_reg);
 			nfnl_nft_nat_set_sreg_proto_min(nle, pmin_reg);
 			nfnl_nft_nat_set_sreg_proto_max(nle, pmax_reg);
-			release_register(ctx);
-			release_register(ctx);
-			break;
+		} else {
+			netlink_gen_expr(ctx, stmt->nat.proto, pmin_reg);
+			nfnl_nft_nat_set_sreg_proto_min(nle, pmin_reg);
 		}
 	}
 
+	while (registers > 0) {
+		release_register(ctx);
+		registers--;
+	}
+
 	nfnl_nft_rule_add_expr(ctx->nlr, nle);
 }