Graph* GetConstantGraph(
    const Graph* orig_graph, const std::vector<Node*>& nodes,
    const std::unordered_map<const Node*, std::vector<Tensor>>&
        shape_replacement_map,
    std::map<NodeAndOutput, NodeAndOutput>* tensors_to_fetch,
    const ConstantFoldNameGenerator& generate_new_name) {
  Graph* constant_graph = new Graph(orig_graph->op_registry());
  std::unordered_map<Node*, std::vector<Node*>> node_map;
  node_map[orig_graph->source_node()] = {constant_graph->source_node()};
  node_map[orig_graph->sink_node()] = {constant_graph->sink_node()};
  for (Node* n : nodes) {
    if (shape_replacement_map.count(n) == 0) {
      AddNodeToConstantGraph(n, &node_map, constant_graph);
    } else {
      AddShapeNodeToConstantGraph(n, shape_replacement_map, &node_map,
                                  generate_new_name, constant_graph);
    }
  }

  for (auto const& added_nodes : node_map) {
    for (const Edge* out_edge : added_nodes.first->out_edges()) {
      if (node_map.count(out_edge->dst()) == 0) {
        if (out_edge->IsControlEdge()) continue;
        if (added_nodes.second.size() == 1) {
          tensors_to_fetch->insert(
              {{added_nodes.second[0], out_edge->src_output()},
               {added_nodes.first, out_edge->src_output()}});
        } else {
          // The node had multiple outputs and was replaced by a
          // vector of constants, so the NodeAndOutput is the 0th
          // output of the kth added constant, rather than the kth
          // output of the added node as in the standard case above.
          tensors_to_fetch->insert(
              {{added_nodes.second[out_edge->src_output()], 0},
               {added_nodes.first, out_edge->src_output()}});
        }
      }
    }
  }

  return constant_graph;
}


void ConsiderConstantFoldableNode(
    Node* n, const ConstantFoldingOptions& opts, std::vector<Node*>* nodes,
    std::unordered_map<const Node*, gtl::FlatSet<Node*>>* constant_control_deps,
    std::unordered_map<const Node*, std::vector<Tensor>>* shape_replacement_map,
    bool* internal_node_inserted) {
  if (IsConstantFoldable(n, opts.shape_map, opts.consider,
                         opts.max_constant_size_in_bytes,
                         shape_replacement_map)) {
    // A node is constant provided all of its non-control incoming Tensors come
    // from constant nodes, or it's a shape Op with statically known inputs in
    // which case it is placed in shape_replacement_map.
    //
    // We allow control dependencies from non-constant nodes to constant nodes,
    // but to preserve the graph structure we must transfer the control
    // dependency onto any constant replacement.
    bool all_parents_constant = true;
    for (const Edge* in : n->in_edges()) {
      // Allows non-constant -> constant control edges.
      if (!in->IsControlEdge() &&
          constant_control_deps->count(in->src()) == 0) {
        all_parents_constant = false;
        break;
      }
    }
    if (all_parents_constant || shape_replacement_map->count(n) != 0) {
      gtl::FlatSet<Node*>& control_deps = (*constant_control_deps)[n];
      for (const Edge* e : n->in_edges()) {
        if (constant_control_deps->count(e->src()) == 0) {
          // This branch is taken if the incoming edge is a control dependency,
          // in which case we want to add it to the dependencies being
          // accumulated for this node, or the incoming edge is not
          // constant. The latter may happen when n is a shape node and the
          // source has known shape. In that case add a control dependency from
          // the source node, since there was previously a data dependency and
          // we want to preserve sequencing constraints.
          if (!e->src()->IsSource()) {
            control_deps.insert(e->src());
          }
        } else {
          // If the parent has been accumulating control dependencies, add all
          // of its transitive control deps.
          const gtl::FlatSet<Node*>& parent_deps =
              (*constant_control_deps)[e->src()];
          control_deps.insert(parent_deps.begin(), parent_deps.end());
        }
      }
      nodes->push_back(n);
      if (!n->IsConstant()) {
        *internal_node_inserted = true;
      }
    }
  }
}

Don't constant-fold DT_RESOURCE constants.

PiperOrigin-RevId: 391803952
Change-Id: I0ea3ec31d3e7dfda0f03b4027a237f08d00a3091


bool IsShapeOp(const Node* n) {
  const auto& ts = n->type_string();
  return ts == "Shape" || ts == "ShapeN" || ts == "Rank" || ts == "Size";
}


void AddNodeToConstantGraph(
    Node* n, std::unordered_map<Node*, std::vector<Node*>>* node_map,
    Graph* constant_graph) {
  std::vector<Node*>& added = (*node_map)[n];
  added.push_back(constant_graph->CopyNode(n));
  for (const Edge* in_edge : n->in_edges()) {
    // Don't copy control edges to the constant graph.
    if (!in_edge->IsControlEdge()) {
      Node* in = in_edge->src();
      auto it = node_map->find(in);
      CHECK(it != node_map->end())
          << n->DebugString() << " <-" << in->DebugString();
      if (it->second.size() == 1) {
        constant_graph->AddEdge(it->second[0], in_edge->src_output(), added[0],
                                in_edge->dst_input());
      } else {
        // The original source node had multiple outputs and was replaced by a
        // vector of constants, so the edge comes from the 0th output of the kth
        // added constant, rather than the kth output of the added node as in
        // the standard case above.
        constant_graph->AddEdge(it->second[in_edge->src_output()], 0, added[0],
                                in_edge->dst_input());
      }
    }
  }
}


bool MaybeReplaceShapeOrShapeNOp(
    const Node* n, const std::vector<PartialTensorShape>& input_shapes,
    std::unordered_map<const Node*, std::vector<Tensor>>*
        shape_replacement_map) {
  std::vector<Tensor> defined_shape;
  for (const auto& shape : input_shapes) {
    if (!shape.IsFullyDefined()) {
      return false;
    }
    const int rank = shape.dims();
    DataType op_type = n->output_type(0);
    Tensor t(op_type, TensorShape({rank}));
    if (op_type == DT_INT64) {
      auto vec = t.vec<int64_t>();
      for (int i = 0; i < rank; ++i) {
        vec(i) = shape.dim_size(i);
      }
    } else {
      CHECK(op_type == DT_INT32);
      auto vec = t.vec<int32>();
      for (int i = 0; i < rank; ++i) {
        if (shape.dim_size(i) > INT_MAX) {
          VLOG(1) << "Node " << n->name() << " has input shape dimension " << i
                  << " of " << shape.dim_size(i) << " but type INT32 "
                  << " so not replacing as constant: this will trigger a "
                     "runtime error later.";
          return false;
        }
        vec(i) = static_cast<int32>(shape.dim_size(i));
      }
    }
    defined_shape.push_back(t);
  }
  // All the inputs had known shapes so we can replace the node by constants
  // later in the rewrite.
  shape_replacement_map->insert({n, defined_shape});
  return true;
}


int64_t UniqueConstantId() {
  static std::atomic_int_fast64_t unique_constant_id;
  return unique_constant_id.fetch_add(1);
}


bool IsConstantFoldable(
    const Node* n,
    const std::unordered_map<string, std::vector<PartialTensorShape>>*
        shape_map,
    const std::function<bool(const Node*)>& consider,
    int64_t max_constant_size_in_bytes,
    std::unordered_map<const Node*, std::vector<Tensor>>*
        shape_replacement_map) {
  if (n->IsConstant()) {
    // Skip constant folding resources as they cannot be deep copied.
    return n->output_type(0) != DT_RESOURCE;
  }
  if (MaybeReplaceShapeOp(n, shape_map, shape_replacement_map)) {
    return true;
  }
  if (n->op_def().is_stateful()) {
    return false;
  }
  if (consider && !consider(n)) {
    return false;
  }
  if (shape_map != nullptr) {
    // We can skip the node if an output is known to be oversized.
    auto shape_it = shape_map->find(n->name());
    if (shape_it != shape_map->end()) {
      for (int64_t i = 0; i < shape_it->second.size(); ++i) {
        const auto& out_shape = shape_it->second[i];
        if (out_shape.IsFullyDefined() &&
            out_shape.num_elements() * DataTypeSize(n->output_type(i)) >
                max_constant_size_in_bytes) {
          return false;
        }
      }
    }
  }
  if (n->IsControlFlow() || n->IsSend() || n->IsRecv()) {
    return false;
  }
  // TODO(yuanbyu): For now disable these session handle operations.
  if (n->IsGetSessionHandle() || n->IsGetSessionTensor() ||
      n->IsDeleteSessionTensor()) {
    return false;
  }
  if (n->IsSource()) {
    return false;
  }
  if (n->IsSink()) {
    return false;
  }
  if (n->IsFakeParam()) {
    return false;
  }
  // Since constant-folding runs on the CPU, do not attempt to constant-fold
  // operators that have no CPU kernel. Also implies that we will not
  // constant-fold functions.
  // TODO(phawkins): allow constant-folding for functions; functions may
  // be arbitrarily expensive to execute.
  if (!KernelDefAvailable(DeviceType(DEVICE_CPU), n->def())) {
    return false;
  }
  // Do not constant fold nodes which will be allocated by ScopedAllocator.
  // This is because the constant-folding graph will not contain the
  // `_ScopedAllocator` node, and that is necessary to be able to run a node
  // that will use this allocator.
  if (n->attrs().Find(kScopedAllocatorAttrName) != nullptr) {
    VLOG(2) << "Skip node [" << n->DebugString()
            << "] for constant folding due to scoped allocator";
    return false;
  }
  return true;
}


Status ConstantFold(const ConstantFoldingOptions& opts,
                    FunctionLibraryRuntime* function_library, Env* env,
                    const Device* partition_device, Graph* graph,
                    bool* was_mutated) {
  // TensorFlow flushes denormals to zero and rounds to nearest, so we do
  // the same here.
  port::ScopedFlushDenormal flush;
  port::ScopedSetRound round(FE_TONEAREST);

  DumpGraph("Before", graph);
  ConstantFoldNameGenerator generate_new_name = opts.generate_new_name;
  if (generate_new_name == nullptr) {
    generate_new_name = [](Graph* graph, string old_name) {
      return strings::StrCat(graph->NewName(old_name), "__cf__",
                             UniqueConstantId());
    };
  }

  std::vector<Node*> constant_foldable_nodes;
  std::unordered_map<const Node*, gtl::FlatSet<Node*>> constant_control_deps;
  std::unordered_map<const Node*, std::vector<Tensor>> shape_replacement_map;
  FindConstantFoldableNodes(graph, opts, &constant_foldable_nodes,
                            &constant_control_deps, &shape_replacement_map);
  if (constant_foldable_nodes.empty()) {
    VLOG(1) << "No constant foldable nodes found";
    *was_mutated = false;
    // This is not an error, so return the status as OK.
    return Status::OK();
  }

  std::map<NodeAndOutput, NodeAndOutput> tensors_to_fetch;
  std::unique_ptr<Graph> constant_graph(
      GetConstantGraph(graph, constant_foldable_nodes, shape_replacement_map,
                       &tensors_to_fetch, generate_new_name));
  DumpGraph("Constant graph", constant_graph.get());

  if (tensors_to_fetch.empty()) {
    VLOG(1) << "No constant nodes found that feed into the original graph.";
    *was_mutated = false;
    // This is not an error, so return the status as OK.
    return Status::OK();
  }
  VLOG(1) << "Constant foldable " << constant_graph->num_node_ids() << " : "
          << graph->num_node_ids();

  std::vector<string> tensors_to_fetch_names;
  std::vector<NodeAndOutput> tensors_to_replace;
  // Sorting the nodes based on the name gives us a stable ordering between runs
  // for the same graph.
  std::vector<std::pair<NodeAndOutput, NodeAndOutput>> tensors_to_fetch_sorted(
      tensors_to_fetch.begin(), tensors_to_fetch.end());
  std::sort(tensors_to_fetch_sorted.begin(), tensors_to_fetch_sorted.end(),
            [](const std::pair<NodeAndOutput, NodeAndOutput>& n1,
               const std::pair<NodeAndOutput, NodeAndOutput>& n2) {
              return std::tie(n1.first.first->name(), n1.first.second) <
                     std::tie(n2.first.first->name(), n2.first.second);
            });
  for (auto n : tensors_to_fetch_sorted) {
    tensors_to_fetch_names.push_back(
        strings::StrCat(n.first.first->name(), ":", n.first.second));
    tensors_to_replace.push_back(n.second);
  }

  auto graph_runner = std::unique_ptr<GraphRunner>(new GraphRunner(env));
  // Evaluate the constant foldable nodes.
  std::vector<Tensor> outputs;
  auto delete_tensors = gtl::MakeCleanup([&graph_runner, &outputs] {
    // Output tensors need to be cleared before the GraphRunner is deleted.
    outputs.clear();
    graph_runner.reset(nullptr);
  });

  Status s =
      graph_runner->Run(constant_graph.get(), function_library, {} /* inputs*/,
                        tensors_to_fetch_names, &outputs);
  if (!s.ok()) {
    VLOG(1) << "Could not fetch constants: " << s;
    *was_mutated = false;
    return s;
  }

  // Fetch the constant tensors and replace the corresponding tensors in the
  // original graph with those constants.
  int32_t num_nodes_replaced = 0;
  for (size_t c = 0; c < outputs.size(); ++c) {
    const gtl::FlatSet<Node*>& control_deps =
        constant_control_deps[tensors_to_replace[c].first];
    if (ReplaceTensorWithConstant(
            graph, partition_device, tensors_to_replace[c], outputs[c],
            control_deps, opts.max_constant_size_in_bytes, generate_new_name)) {
      ++num_nodes_replaced;
    }
  }

  DumpGraph("After", graph);

  *was_mutated = (num_nodes_replaced > 0);
  return Status::OK();
}


bool MaybeReplaceSizeOp(const Node* n,
                        const std::vector<PartialTensorShape>& input_shapes,
                        std::unordered_map<const Node*, std::vector<Tensor>>*
                            shape_replacement_map) {
  CHECK_EQ(input_shapes.size(), 1);
  if (!input_shapes[0].IsFullyDefined()) {
    return false;
  }
  DataType op_type = n->output_type(0);
  Tensor t(op_type, TensorShape({}));
  int64_t size = input_shapes[0].num_elements();
  if (op_type == DT_INT64) {
    t.scalar<int64_t>()() = size;
  } else {
    CHECK(op_type == DT_INT32);
    if (size > INT_MAX) {
      VLOG(1) << "Node " << n->name() << " has input shape size " << size
              << " but type INT32 "
              << " so not replacing as constant: this will trigger a runtime "
                 "error later.";
      return false;
    }
    t.scalar<int32>()() = static_cast<int32>(size);
  }
  shape_replacement_map->insert({n, {t}});
  return true;
}


void FindConstantFoldableNodes(
    const Graph* graph, const ConstantFoldingOptions& opts,
    std::vector<Node*>* nodes,
    std::unordered_map<const Node*, gtl::FlatSet<Node*>>* constant_control_deps,
    std::unordered_map<const Node*, std::vector<Tensor>>*
        shape_replacement_map) {
  bool internal_node_inserted = false;
  // Walk the nodes in data flow order.
  ReverseDFS(
      *graph, nullptr,
      [nodes, constant_control_deps, shape_replacement_map,
       &internal_node_inserted, &opts](Node* n) {
        ConsiderConstantFoldableNode(n, opts, nodes, constant_control_deps,
                                     shape_replacement_map,
                                     &internal_node_inserted);
      },
      NodeComparatorName());
  // If we have inserted just leaf level nodes, then there is nothing to fold.
  if (!internal_node_inserted) {
    nodes->clear();
    constant_control_deps->clear();
  }
}


bool MaybeReplaceRankOp(const Node* n,
                        const std::vector<PartialTensorShape>& input_shapes,
                        std::unordered_map<const Node*, std::vector<Tensor>>*
                            shape_replacement_map) {
  CHECK_EQ(input_shapes.size(), 1);
  if (input_shapes[0].unknown_rank()) {
    return false;
  }
  Tensor t(DT_INT32, TensorShape({}));
  t.scalar<int32>()() = input_shapes[0].dims();
  shape_replacement_map->insert({n, {t}});
  return true;
}


R_API bool r_io_bank_map_add_top(RIO *io, const ut32 bankid, const ut32 mapid) {
	RIOBank *bank = r_io_bank_get (io, bankid);
	RIOMap *map = r_io_map_get (io, mapid);
	r_return_val_if_fail (io && bank && map, false);
	RIOMapRef *mapref = _mapref_from_map (map);
	if (!mapref) {
		return false;
	}
	RIOSubMap *sm = r_io_submap_new (io, mapref);
	if (!sm) {
		free (mapref);
		return false;
	}
	RRBNode *entry = _find_entry_submap_node (bank, sm);
	if (!entry) {
		// no intersection with any submap, so just insert
		if (!r_crbtree_insert (bank->submaps, sm, _find_sm_by_from_vaddr_cb, NULL)) {
			free (sm);
			free (mapref);
			return false;
		}
		r_list_append (bank->maprefs, mapref);
		return true;
	}
	bank->last_used = NULL;
	RIOSubMap *bd = (RIOSubMap *)entry->data;
	if (r_io_submap_to (bd) == r_io_submap_to (sm) &&
		r_io_submap_from (bd) >= r_io_submap_from (sm)) {
		// _find_entry_submap_node guarantees, that there is no submap
		// prior to bd in the range of sm, so instead of deleting and inserting
		// we can just memcpy
		memcpy (bd, sm, sizeof (RIOSubMap));
		free (sm);
		r_list_append (bank->maprefs, mapref);
		return true;
	}
	if (r_io_submap_from (bd) < r_io_submap_from (sm) &&
		r_io_submap_to (sm) < r_io_submap_to (bd)) {
		// split bd into 2 maps => bd and bdsm
		RIOSubMap *bdsm = R_NEWCOPY (RIOSubMap, bd);
		if (!bdsm) {
			free (sm);
			free (mapref);
			return false;
		}
		r_io_submap_set_from (bdsm, r_io_submap_to (sm) + 1);
		r_io_submap_set_to (bd, r_io_submap_from (sm) - 1);
		// TODO: insert and check return value, before adjusting sm size
		if (!r_crbtree_insert (bank->submaps, sm, _find_sm_by_from_vaddr_cb, NULL)) {
			free (sm);
			free (bdsm);
			free (mapref);
			return false;
		}
		if (!r_crbtree_insert (bank->submaps, bdsm, _find_sm_by_from_vaddr_cb, NULL)) {
			r_crbtree_delete (bank->submaps, sm, _find_sm_by_from_vaddr_cb, NULL);
			free (sm);
			free (bdsm);
			free (mapref);
			return false;
		}
		r_list_append (bank->maprefs, mapref);
		return true;
	}

	// guaranteed intersection
	if (r_io_submap_from (bd) < r_io_submap_from (sm)) {
		r_io_submap_set_to (bd, r_io_submap_from (sm) - 1);
		entry = r_rbnode_next (entry);
	}
	ut64 smto = r_io_submap_to (sm);
	if (entry) {
		ut64 ento = r_io_submap_to (((RIOSubMap*)entry->data));
		while (entry && ento <= smto) {
			//delete all submaps that are completly included in sm
			RRBNode *next = r_rbnode_next (entry);
			// this can be optimized, there is no need to do search here
			// XXX this is a workaround to avoid an UAF in Reproducer: iobank-crash
			ento = r_io_submap_to (((RIOSubMap*)entry->data));
			bool a = r_crbtree_delete (bank->submaps, entry->data, _find_sm_by_from_vaddr_cb, NULL);
			if (!a) {
				next = NULL;
			}
			entry = next;
		}
	}
	if (entry && r_io_submap_from (((RIOSubMap *)entry->data)) <= r_io_submap_to (sm)) {
		bd = (RIOSubMap *)entry->data;
		r_io_submap_set_from (bd, r_io_submap_to (sm) + 1);
	}
	if (!r_crbtree_insert (bank->submaps, sm, _find_sm_by_from_vaddr_cb, NULL)) {
		free (sm);
		free (mapref);
		return false;
	}
	r_list_append (bank->maprefs, mapref);
	return true;
}

Properly fix the UAF in r_io_bank_map_add_top ##crash

* Associated with the CVE-2022-0559
* Reported by alkyne Choi via huntr.dev


stradd(char *d)
{
#ifdef MULTIBYTE_SUPPORT
    char *ums, *ups;
    int upslen, eol = 0;
    mbstate_t mbs;

    memset(&mbs, 0, sizeof mbs);
    ums = ztrdup(d);
    ups = unmetafy(ums, &upslen);

    /*
     * We now have a raw string of possibly multibyte characters.
     * Read each character one by one.
     */
    while (upslen > 0) {
	wchar_t cc;
	char *pc;
	size_t cnt = eol ? MB_INVALID : mbrtowc(&cc, ups, upslen, &mbs);

	switch (cnt) {
	case MB_INCOMPLETE:
	    eol = 1;
	    /* FALL THROUGH */
	case MB_INVALID:
	    /* Bad character.  Take the next byte on its own. */
	    pc = nicechar(*ups);
	    cnt = 1;
	    memset(&mbs, 0, sizeof mbs);
	    break;
	case 0:
	    cnt = 1;
	    /* FALL THROUGH */
	default:
	    /* Take full wide character in one go */
	    mb_charinit();
	    pc = wcs_nicechar(cc, NULL, NULL);
	    break;
	}
	/* Keep output as metafied string. */
	addbufspc(strlen(pc));

	upslen -= cnt;
	ups += cnt;

	/* Put printed representation into the buffer */
	while (*pc)
	    *bv->bp++ = *pc++;
    }

    free(ums);
#else
    char *ps, *pc;
    addbufspc(niceztrlen(d));
    /* This loop puts the nice representation of the string into the
     * prompt buffer. */
    for (ps = d; *ps; ps++) {
	for (pc = nicechar(*ps == Meta ? *++ps^32 : *ps); *pc; pc++)
	    *bv->bp++ = *pc;
    }
#endif
}

security/41: Don't perform PROMPT_SUBST evaluation on %F/%K arguments

Mitigates CVE-2021-45444


putstr(int d)
{
    addbufspc(1);
    pputc(d);
    return 0;
}


allocate_colour_buffer(void)
{
    char **atrs;
    int lenfg, lenbg, len;

    if (colseq_buf_allocs++)
	return;

    atrs = getaparam("zle_highlight");
    if (atrs) {
	for (; *atrs; atrs++) {
	    if (strpfx("fg_start_code:", *atrs)) {
		set_colour_code(*atrs + 14, &fg_bg_sequences[COL_SEQ_FG].start);
	    } else if (strpfx("fg_default_code:", *atrs)) {
		set_colour_code(*atrs + 16, &fg_bg_sequences[COL_SEQ_FG].def);
	    } else if (strpfx("fg_end_code:", *atrs)) {
		set_colour_code(*atrs + 12, &fg_bg_sequences[COL_SEQ_FG].end);
	    } else if (strpfx("bg_start_code:", *atrs)) {
		set_colour_code(*atrs + 14, &fg_bg_sequences[COL_SEQ_BG].start);
	    } else if (strpfx("bg_default_code:", *atrs)) {
		set_colour_code(*atrs + 16, &fg_bg_sequences[COL_SEQ_BG].def);
	    } else if (strpfx("bg_end_code:", *atrs)) {
		set_colour_code(*atrs + 12, &fg_bg_sequences[COL_SEQ_BG].end);
	    }
	}
    }

    lenfg = strlen(fg_bg_sequences[COL_SEQ_FG].def);
    /* always need 1 character for non-default code */
    if (lenfg < 1)
	lenfg = 1;
    lenfg += strlen(fg_bg_sequences[COL_SEQ_FG].start) +
	strlen(fg_bg_sequences[COL_SEQ_FG].end);

    lenbg = strlen(fg_bg_sequences[COL_SEQ_BG].def);
    /* always need 1 character for non-default code */
    if (lenbg < 1)
	lenbg = 1;
    lenbg += strlen(fg_bg_sequences[COL_SEQ_BG].start) +
	strlen(fg_bg_sequences[COL_SEQ_BG].end);

    len = lenfg > lenbg ? lenfg : lenbg;
    /* add 1 for the null and 14 for truecolor */
    colseq_buf = (char *)zalloc(len+15);
}


free_colour_buffer(void)
{
    if (--colseq_buf_allocs)
	return;

    DPUTS(!colseq_buf, "Freeing colour sequence buffer without alloc");
    /* Free buffer for colour code composition */
    free(colseq_buf);
    colseq_buf = NULL;
}


match_highlight(const char *teststr, zattr *on_var)
{
    int found = 1;

    *on_var = 0;
    while (found && *teststr) {
	const struct highlight *hl;

	found = 0;
	if (strpfx("fg=", teststr) || strpfx("bg=", teststr)) {
	    int is_fg = (teststr[0] == 'f');
	    zattr atr;

	    teststr += 3;
	    atr = match_colour(&teststr, is_fg, 0);
	    if (*teststr == ',')
		teststr++;
	    else if (*teststr)
		break;
	    found = 1;
	    /* skip out of range colours but keep scanning attributes */
	    if (atr != TXT_ERROR)
		*on_var |= atr;
	} else {
	    for (hl = highlights; hl->name; hl++) {
		if (strpfx(hl->name, teststr)) {
		    const char *val = teststr + strlen(hl->name);

		    if (*val == ',')
			val++;
		    else if (*val)
			break;

		    *on_var |= hl->mask_on;
		    *on_var &= ~hl->mask_off;
		    teststr = val;
		    found = 1;
		}
	    }
	}
    }
}


pputc(char c)
{
    if (imeta(c)) {
	*bv->bp++ = Meta;
	c ^= 32;
    }
    *bv->bp++ = c;
    if (c == '\n' && !bv->dontcount)
	bv->bufline = bv->bp;
}


set_colour_code(char *str, char **var)
{
    char *keyseq;
    int len;

    zsfree(*var);
    keyseq = getkeystring(str, &len, GETKEYS_BINDKEY, NULL);
    *var = metafy(keyseq, len, META_DUP);
}


set_default_colour_sequences(void)
{
    fg_bg_sequences[COL_SEQ_FG].start = ztrdup(TC_COL_FG_START);
    fg_bg_sequences[COL_SEQ_FG].end = ztrdup(TC_COL_FG_END);
    fg_bg_sequences[COL_SEQ_FG].def = ztrdup(TC_COL_FG_DEFAULT);

    fg_bg_sequences[COL_SEQ_BG].start = ztrdup(TC_COL_BG_START);
    fg_bg_sequences[COL_SEQ_BG].end = ztrdup(TC_COL_BG_END);
    fg_bg_sequences[COL_SEQ_BG].def = ztrdup(TC_COL_BG_DEFAULT);
}


match_named_colour(const char **teststrp)
{
    const char *teststr = *teststrp, *end, **cptr;
    int len;

    for (end = teststr; ialpha(*end); end++)
	;
    len = end - teststr;
    *teststrp = end;

    for (cptr = ansi_colours; *cptr; cptr++) {
	if (!strncmp(teststr, *cptr, len))
	    return cptr - ansi_colours;
    }

    return -1;
}


output_highlight(zattr atr, char *buf)
{
    const struct highlight *hp;
    int atrlen = 0, len;
    char *ptr = buf;

    if (atr & TXTFGCOLOUR) {
	len = output_colour(txtchangeget(atr, TXT_ATTR_FG_COL),
			    COL_SEQ_FG,
			    (atr & TXT_ATTR_FG_TERMCAP),
			    (atr & TXT_ATTR_FG_24BIT),
			    ptr);
	atrlen += len;
	if (buf)
	    ptr += len;
    }
    if (atr & TXTBGCOLOUR) {
	if (atrlen) {
	    atrlen++;
	    if (buf) {
		strcpy(ptr, ",");
		ptr++;
	    }
	}
	len = output_colour(txtchangeget(atr, TXT_ATTR_BG_COL),
			    COL_SEQ_BG,
			    (atr & TXT_ATTR_BG_TERMCAP),
			    (atr & TXT_ATTR_BG_24BIT),
			    ptr);
	atrlen += len;
	if (buf)
	    ptr += len;
    }
    for (hp = highlights; hp->name; hp++) {
	if (hp->mask_on & atr) {
	    if (atrlen) {
		atrlen++;
		if (buf) {
		    strcpy(ptr, ",");
		    ptr++;
		}
	    }
	    len = strlen(hp->name);
	    atrlen += len;
	    if (buf) {
		strcpy(ptr, hp->name);
		ptr += len;
	    }
	}
    }

    if (atrlen == 0) {
	if (buf)
	    strcpy(ptr, "none");
	return 4;
    }
    return atrlen;
}


output_colour(int colour, int fg_bg, int use_tc, int truecol, char *buf)
{
    int atrlen = 3, len;
    char *ptr = buf;
    if (buf) {
	strcpy(ptr, fg_bg == COL_SEQ_FG ? "fg=" : "bg=");
	ptr += 3;
    }
    if (truecol) {
	/* length of hex triplet always 7, don't need sprintf to count */
	atrlen += buf ? sprintf(ptr, "#%02x%02x%02x", colour >> 16,
		(colour >> 8) & 0xff, colour & 0xff) : 7;
    /* colour should only be > 7 if using termcap but let's be safe. Update:
     * currently other places in code don't always imply that colour > 7 =>
     * using-termcap - if zle_highlight will be non-default, then it will be
     * used instead of termcap even for colour > 7. Here this just emits the
     * color number, so it works fine for both zle_highlight and tercap cases
     */
    } else if (use_tc || colour > 7) {
	char digbuf[DIGBUFSIZE];
	sprintf(digbuf, "%d", colour);
	len = strlen(digbuf);
	atrlen += len;
	if (buf)
	    strcpy(ptr, digbuf);
    } else {
	len = strlen(ansi_colours[colour]);
	atrlen += len;
	if (buf)
	    strcpy(ptr, ansi_colours[colour]);
    }

    return atrlen;
}


countprompt(char *str, int *wp, int *hp, int overf)
{
    int w = 0, h = 1, multi = 0, wcw = 0;
    int s = 1;
#ifdef MULTIBYTE_SUPPORT
    char inchar;
    mbstate_t mbs;
    wchar_t wc;

    memset(&mbs, 0, sizeof(mbs));
#endif

    for (; *str; str++) {
	/*
	 * Avoid double-incrementing the height when there's a newline in the
	 * prompt and the line it terminates takes up exactly the width of the
	 * terminal
	 */
	while (w > zterm_columns && overf >= 0 && !multi) {
	    h++;
	    if (wcw) {
		/*
		 * Wide characters don't get split off. They move to the
		 * next line if there is not enough space.
		 */
		w = wcw;
		break;
	    } else {
		/*
		 * Tabs overflow to the next line as if they were made of spaces.
		 */
		w -= zterm_columns;
	    }
	}
	wcw = 0;
	/*
	 * Input string should be metafied, so tokens in it should
	 * be real tokens, even if there are multibyte characters.
	 */
	if (*str == Inpar)
	    s = 0;
	else if (*str == Outpar)
	    s = 1;
	else if (*str == Nularg)
	    w++;
	else if (s) {
	    if (*str == Meta) {
#ifdef MULTIBYTE_SUPPORT
		inchar = *++str ^ 32;
#else
		str++;
#endif
	    } else {
#ifdef MULTIBYTE_SUPPORT
		/*
		 * Don't look for tab or newline in the middle
		 * of a multibyte character.  Otherwise, we are
		 * relying on the character set being an extension
		 * of ASCII so it's safe to test a single byte.
		 */
		if (!multi) {
#endif
		    if (*str == '\t') {
			w = (w | 7) + 1;
			continue;
		    } else if (*str == '\n') {
			w = 0;
			h++;
			continue;
		    }
#ifdef MULTIBYTE_SUPPORT
		}

		inchar = *str;
#endif
	    }

#ifdef MULTIBYTE_SUPPORT
	    switch (mbrtowc(&wc, &inchar, 1, &mbs)) {
	    case MB_INCOMPLETE:
		/* Character is incomplete -- keep looking. */
		multi = 1;
		break;
	    case MB_INVALID:
		memset(&mbs, 0, sizeof mbs);
		/* Invalid character: assume single width. */
		multi = 0;
		w++;
		break;
	    case 0:
		multi = 0;
		break;
	    default:
		/*
		 * If the character isn't printable, WCWIDTH() returns
		 * -1.  We assume width 1.
		 */
		wcw = WCWIDTH(wc);
		if (wcw >= 0)
		    w += wcw;
		else
		    w++;
		multi = 0;
		break;
	    }
#else
	    w++;
#endif
	}
    }
    /*
     * multi may still be set if we were in the middle of the character.
     * This isn't easy to handle generally; just assume there's no
     * output.
     */
    while (w > zterm_columns && overf >= 0) {
	h++;
	if (wcw) {
	    w = wcw;
	    break;
	} else {
	    w -= zterm_columns;
	}
    }
    if (w == zterm_columns && overf == 0) {
	w = 0;
	h++;
    }
    if(wp)
	*wp = w;
    if(hp)
	*hp = h;
}


promptpath(char *p, int npath, int tilde)
{
    char *modp = p;
    Nameddir nd;

    if (tilde && ((nd = finddir(p))))
	modp = tricat("~", nd->node.nam, p + strlen(nd->dir));

    if (npath) {
	char *sptr;
	if (npath > 0) {
	    for (sptr = modp + strlen(modp); sptr > modp; sptr--) {
		if (*sptr == '/' && !--npath) {
		    sptr++;
		    break;
		}
	    }
	    if (*sptr == '/' && sptr[1] && sptr != modp)
		sptr++;
	    stradd(sptr);
	} else {
	    char cbu;
	    for (sptr = modp+1; *sptr; sptr++)
		if (*sptr == '/' && !++npath)
		    break;
	    cbu = *sptr;
	    *sptr = 0;
	    stradd(modp);
	    *sptr = cbu;
	}
    } else
	stradd(modp);

    if (p != modp)
	zsfree(modp);
}


match_colour(const char **teststrp, int is_fg, int colour)
{
    int shft, named = 0, tc;
    zattr on;

    if (is_fg) {
	shft = TXT_ATTR_FG_COL_SHIFT;
	on = TXTFGCOLOUR;
	tc = TCFGCOLOUR;
    } else {
	shft = TXT_ATTR_BG_COL_SHIFT;
	on = TXTBGCOLOUR;
	tc = TCBGCOLOUR;
    }
    if (teststrp) {
	if (**teststrp == '#' && isxdigit((*teststrp)[1])) {
	    struct color_rgb color;
	    char *end;
	    zlong col = zstrtol(*teststrp+1, &end, 16);
            if (end - *teststrp == 4) {
		color.red = col >> 8 | ((col >> 8) << 4);
		color.green = (col & 0xf0) >> 4;
		color.green |= color.green << 4;
		color.blue = col & 0xf;
		color.blue |= color.blue << 4;
	    } else if (end - *teststrp == 7) {
		color.red = col >> 16;
		color.green = (col & 0xff00) >> 8;
		color.blue = col & 0xff;
	    } else
		return TXT_ERROR;
	    *teststrp = end;
	    colour = runhookdef(GETCOLORATTR, &color) - 1;
	    if (colour == -1) { /* no hook function added, try true color (24-bit) */
		colour = (((color.red << 8) + color.green) << 8) + color.blue;
		return on | (is_fg ? TXT_ATTR_FG_24BIT : TXT_ATTR_BG_24BIT) |
			(zattr)colour << shft;
	    } else if (colour <= -2) {
		return TXT_ERROR;
	    }
	} else if ((named = ialpha(**teststrp))) {
	    colour = match_named_colour(teststrp);
	    if (colour == 8) {
		/* default */
		return is_fg ? TXTNOFGCOLOUR : TXTNOBGCOLOUR;
	    }
	    if (colour < 0)
		return TXT_ERROR;
	}
	else {
	    colour = (int)zstrtol(*teststrp, (char **)teststrp, 10);
	    if (colour < 0 || colour >= 256)
		return TXT_ERROR;
	}
    }
    /*
     * Try termcap for numbered characters if possible.
     * Don't for named characters, since our best bet
     * of getting the names right is with ANSI sequences.
     */
    if (!named && tccan(tc)) {
	if (tccolours >= 0 && colour >= tccolours) {
	    /*
	     * Out of range of termcap colours.
	     * Can we assume ANSI colours work?
	     */
	    if (colour > 7)
		return TXT_ERROR; /* No. */
	} else {
	    /*
	     * We can handle termcap colours and the number
	     * is in range, so use termcap.
	     */
	    on |= is_fg ? TXT_ATTR_FG_TERMCAP :
		TXT_ATTR_BG_TERMCAP;
	}
    }
    return on | (zattr)colour << shft;
}


parsecolorchar(zattr arg, int is_fg)
{
    if (bv->fm[1] == '{') {
	char *ep;
	bv->fm += 2; /* skip over F{ */
	if ((ep = strchr(bv->fm, '}'))) {
	    char oc = *ep, *col, *coll;
	    int ops = opts[PROMPTSUBST], opb = opts[PROMPTBANG];
	    int opp = opts[PROMPTPERCENT];

	    opts[PROMPTPERCENT] = 1;
	    opts[PROMPTSUBST] = opts[PROMPTBANG] = 0;

	    *ep = '\0';
	    /* expand the contents of the argument so you can use
	     * %v for example */
	    coll = col = promptexpand(bv->fm, 0, NULL, NULL, NULL);
	    *ep = oc;
	    arg = match_colour((const char **)&coll, is_fg, 0);
	    free(col);
	    bv->fm = ep;

	    opts[PROMPTSUBST] = ops;
	    opts[PROMPTBANG] = opb;
	    opts[PROMPTPERCENT] = opp;
	} else {
	    arg = match_colour((const char **)&bv->fm, is_fg, 0);
	    if (*bv->fm != '}')
		bv->fm--;
	}
    } else
	arg = match_colour(NULL, 1, arg);
    return arg;
}


putpromptchar(int doprint, int endchar, zattr *txtchangep)
{
    char *ss, *hostnam;
    int t0, arg, test, sep, j, numjobs, len;
    zattr atr;
    struct tm *tm;
    struct timespec ts;
    time_t timet;
    Nameddir nd;

    for (; *bv->fm && *bv->fm != endchar; bv->fm++) {
	arg = 0;
	if (*bv->fm == '%' && isset(PROMPTPERCENT)) {
	    int minus = 0;
	    bv->fm++;
	    if (*bv->fm == '-') {
		minus = 1;
		bv->fm++;
	    }
	    if (idigit(*bv->fm)) {
		arg = zstrtol(bv->fm, &bv->fm, 10);
		if (minus)
		    arg *= -1;
	    } else if (minus)
		arg = -1;
	    if (*bv->fm == '(') {
		int tc, otruncwidth;

		if (idigit(*++bv->fm)) {
		    arg = zstrtol(bv->fm, &bv->fm, 10);
		} else if (arg < 0) {
		    /* negative numbers don't make sense here */
		    arg *= -1;
		}
		test = 0;
		ss = pwd;
		switch (tc = *bv->fm) {
		case 'c':
		case '.':
		case '~':
		    if ((nd = finddir(ss))) {
			arg--;
			ss += strlen(nd->dir);
		    } /*FALLTHROUGH*/
		case '/':
		case 'C':
		    /* `/' gives 0, `/any' gives 1, etc. */
		    if (*ss && *ss++ == '/' && *ss)
			arg--;
		    for (; *ss; ss++)
			if (*ss == '/')
			    arg--;
		    if (arg <= 0)
			test = 1;
		    break;
		case 't':
		case 'T':
		case 'd':
		case 'D':
		case 'w':
		    timet = time(NULL);
		    tm = localtime(&timet);
		    switch (tc) {
		    case 't':
			test = (arg == tm->tm_min);
			break;
		    case 'T':
			test = (arg == tm->tm_hour);
			break;
		    case 'd':
			test = (arg == tm->tm_mday);
			break;
		    case 'D':
			test = (arg == tm->tm_mon);
			break;
		    case 'w':
			test = (arg == tm->tm_wday);
			break;
		    }
		    break;
		case '?':
		    if (lastval == arg)
			test = 1;
		    break;
		case '#':
		    if (geteuid() == (uid_t)arg)
			test = 1;
		    break;
		case 'g':
		    if (getegid() == (gid_t)arg)
			test = 1;
		    break;
		case 'j':
		    for (numjobs = 0, j = 1; j <= maxjob; j++)
			if (jobtab[j].stat && jobtab[j].procs &&
		    	    !(jobtab[j].stat & STAT_NOPRINT)) numjobs++;
		    if (numjobs >= arg)
		    	test = 1;
		    break;
		case 'l':
		    *bv->bp = '\0';
		    countprompt(bv->bufline, &t0, 0, 0);
		    if (minus)
			t0 = zterm_columns - t0;
		    if (t0 >= arg)
			test = 1;
		    break;
		case 'e':
		    {
			Funcstack fsptr = funcstack;
			test = arg;
			while (fsptr && test > 0) {
			    test--;
			    fsptr = fsptr->prev;
			}
			test = !test;
		    }
		    break;
		case 'L':
		    if (shlvl >= arg)
			test = 1;
		    break;
		case 'S':
		    if (time(NULL) - shtimer.tv_sec >= arg)
			test = 1;
		    break;
		case 'v':
		    if (arrlen_ge(psvar, arg))
			test = 1;
		    break;
		case 'V':
		    if (psvar && *psvar && arrlen_ge(psvar, arg)) {
			if (*psvar[(arg ? arg : 1) - 1])
			    test = 1;
		    }
		    break;
		case '_':
		    test = (cmdsp >= arg);
		    break;
		case '!':
		    test = privasserted();
		    break;
		default:
		    test = -1;
		    break;
		}
		if (!*bv->fm || !(sep = *++bv->fm))
		    return 0;
		bv->fm++;
		/* Don't do the current truncation until we get back */
		otruncwidth = bv->truncwidth;
		bv->truncwidth = 0;
		if (!putpromptchar(test == 1 && doprint, sep,
				   txtchangep) || !*++bv->fm ||
		    !putpromptchar(test == 0 && doprint, ')',
				   txtchangep)) {
		    bv->truncwidth = otruncwidth;
		    return 0;
		}
		bv->truncwidth = otruncwidth;
		continue;
	    }
	    if (!doprint)
		switch(*bv->fm) {
		  case '[':
		    while(idigit(*++bv->fm));
		    while(*++bv->fm != ']');
		    continue;
		  case '<':
		    while(*++bv->fm != '<');
		    continue;
		  case '>':
		    while(*++bv->fm != '>');
		    continue;
		  case 'D':
		    if(bv->fm[1]=='{')
			while(*++bv->fm != '}');
		    continue;
		  default:
		    continue;
		}
	    switch (*bv->fm) {
	    case '~':
		promptpath(pwd, arg, 1);
		break;
	    case 'd':
	    case '/':
		promptpath(pwd, arg, 0);
		break;
	    case 'c':
	    case '.':
		promptpath(pwd, arg ? arg : 1, 1);
		break;
	    case 'C':
		promptpath(pwd, arg ? arg : 1, 0);
		break;
	    case 'N':
		promptpath(scriptname ? scriptname : argzero, arg, 0);
		break;
	    case 'h':
	    case '!':
		addbufspc(DIGBUFSIZE);
		convbase(bv->bp, curhist, 10);
		bv->bp += strlen(bv->bp);
		break;
	    case 'j':
		for (numjobs = 0, j = 1; j <= maxjob; j++)
		    if (jobtab[j].stat && jobtab[j].procs &&
		    	!(jobtab[j].stat & STAT_NOPRINT)) numjobs++;
		addbufspc(DIGBUFSIZE);
		sprintf(bv->bp, "%d", numjobs);
		bv->bp += strlen(bv->bp);
		break;
	    case 'M':
		queue_signals();
		if ((hostnam = getsparam("HOST")))
		    stradd(hostnam);
		unqueue_signals();
		break;
	    case 'm':
		if (!arg)
		    arg++;
		queue_signals();
		if (!(hostnam = getsparam("HOST"))) {
		    unqueue_signals();
		    break;
		}
		if (arg < 0) {
		    for (ss = hostnam + strlen(hostnam); ss > hostnam; ss--)
			if (ss[-1] == '.' && !++arg)
			    break;
		    stradd(ss);
		} else {
		    for (ss = hostnam; *ss; ss++)
			if (*ss == '.' && !--arg)
			    break;
		    stradd(*ss ? dupstrpfx(hostnam, ss - hostnam) : hostnam);
		}
		unqueue_signals();
		break;
	    case 'S':
		txtchangeset(txtchangep, TXTSTANDOUT, TXTNOSTANDOUT);
		txtset(TXTSTANDOUT);
		tsetcap(TCSTANDOUTBEG, TSC_PROMPT);
		break;
	    case 's':
		txtchangeset(txtchangep, TXTNOSTANDOUT, TXTSTANDOUT);
		txtunset(TXTSTANDOUT);
		tsetcap(TCSTANDOUTEND, TSC_PROMPT|TSC_DIRTY);
		break;
	    case 'B':
		txtchangeset(txtchangep, TXTBOLDFACE, TXTNOBOLDFACE);
		txtset(TXTBOLDFACE);
		tsetcap(TCBOLDFACEBEG, TSC_PROMPT|TSC_DIRTY);
		break;
	    case 'b':
		txtchangeset(txtchangep, TXTNOBOLDFACE, TXTBOLDFACE);
		txtunset(TXTBOLDFACE);
		tsetcap(TCALLATTRSOFF, TSC_PROMPT|TSC_DIRTY);
		break;
	    case 'U':
		txtchangeset(txtchangep, TXTUNDERLINE, TXTNOUNDERLINE);
		txtset(TXTUNDERLINE);
		tsetcap(TCUNDERLINEBEG, TSC_PROMPT);
		break;
	    case 'u':
		txtchangeset(txtchangep, TXTNOUNDERLINE, TXTUNDERLINE);
		txtunset(TXTUNDERLINE);
		tsetcap(TCUNDERLINEEND, TSC_PROMPT|TSC_DIRTY);
		break;
	    case 'F':
		atr = parsecolorchar(arg, 1);
		if (!(atr & (TXT_ERROR | TXTNOFGCOLOUR))) {
		    txtchangeset(txtchangep, atr & TXT_ATTR_FG_ON_MASK,
				 TXTNOFGCOLOUR | TXT_ATTR_FG_COL_MASK);
		    txtunset(TXT_ATTR_FG_COL_MASK);
		    txtset(atr & TXT_ATTR_FG_ON_MASK);
		    set_colour_attribute(atr, COL_SEQ_FG, TSC_PROMPT);
		    break;
		}
		/* else FALLTHROUGH */
	    case 'f':
		txtchangeset(txtchangep, TXTNOFGCOLOUR, TXT_ATTR_FG_ON_MASK);
		txtunset(TXT_ATTR_FG_ON_MASK);
		set_colour_attribute(TXTNOFGCOLOUR, COL_SEQ_FG, TSC_PROMPT);
		break;
	    case 'K':
		atr = parsecolorchar(arg, 0);
		if (!(atr & (TXT_ERROR | TXTNOBGCOLOUR))) {
		    txtchangeset(txtchangep, atr & TXT_ATTR_BG_ON_MASK,
				 TXTNOBGCOLOUR | TXT_ATTR_BG_COL_MASK);
		    txtunset(TXT_ATTR_BG_COL_MASK);
		    txtset(atr & TXT_ATTR_BG_ON_MASK);
		    set_colour_attribute(atr, COL_SEQ_BG, TSC_PROMPT);
		    break;
		}
		/* else FALLTHROUGH */
	    case 'k':
		txtchangeset(txtchangep, TXTNOBGCOLOUR, TXT_ATTR_BG_ON_MASK);
		txtunset(TXT_ATTR_BG_ON_MASK);
		set_colour_attribute(TXTNOBGCOLOUR, COL_SEQ_BG, TSC_PROMPT);
		break;
	    case '[':
		if (idigit(*++bv->fm))
		    arg = zstrtol(bv->fm, &bv->fm, 10);
		if (!prompttrunc(arg, ']', doprint, endchar, txtchangep))
		    return *bv->fm;
		break;
	    case '<':
	    case '>':
		/* Test (minus) here so -0 means "at the right margin" */
		if (minus) {
		    *bv->bp = '\0';
		    countprompt(bv->bufline, &t0, 0, 0);
		    arg = zterm_columns - t0 + arg;
		    if (arg <= 0)
			arg = 1;
		}
		if (!prompttrunc(arg, *bv->fm, doprint, endchar, txtchangep))
		    return *bv->fm;
		break;
	    case '{': /*}*/
		if (!bv->dontcount++) {
		    addbufspc(1);
		    *bv->bp++ = Inpar;
		}
		if (arg <= 0)
		    break;
		/* else */
		/* FALLTHROUGH */
	    case 'G':
		if (arg > 0) {
		    addbufspc(arg);
		    while (arg--)
			*bv->bp++ = Nularg;
		} else {
		    addbufspc(1);
		    *bv->bp++ = Nularg;
		}
		break;
	    case /*{*/ '}':
		if (bv->trunccount && bv->trunccount >= bv->dontcount)
		    return *bv->fm;
		if (bv->dontcount && !--bv->dontcount) {
		    addbufspc(1);
		    *bv->bp++ = Outpar;
		}
		break;
	    case 't':
	    case '@':
	    case 'T':
	    case '*':
	    case 'w':
	    case 'W':
	    case 'D':
		{
		    char *tmfmt, *dd, *tmbuf = NULL;

		    switch (*bv->fm) {
		    case 'T':
			tmfmt = "%K:%M";
			break;
		    case '*':
			tmfmt = "%K:%M:%S";
			break;
		    case 'w':
			tmfmt = "%a %f";
			break;
		    case 'W':
			tmfmt = "%m/%d/%y";
			break;
		    case 'D':
			if (bv->fm[1] == '{' /*}*/) {
			    for (ss = bv->fm + 2; *ss && *ss != /*{*/ '}'; ss++)
				if(*ss == '\\' && ss[1])
				    ss++;
			    dd = tmfmt = tmbuf = zalloc(ss - bv->fm);
			    for (ss = bv->fm + 2; *ss && *ss != /*{*/ '}';
				 ss++) {
				if(*ss == '\\' && ss[1])
				    ss++;
				*dd++ = *ss;
			    }
			    *dd = 0;
			    bv->fm = ss - !*ss;
			    if (!*tmfmt) {
				free(tmbuf);
				continue;
			    }
			} else
			    tmfmt = "%y-%m-%d";
			break;
		    default:
			tmfmt = "%l:%M%p";
			break;
		    }
		    zgettime(&ts);
		    tm = localtime(&ts.tv_sec);
		    /*
		     * Hack because strftime won't say how
		     * much space it actually needs.  Try to add it
		     * a few times until it works.  Some formats don't
		     * actually have a length, so we could go on for
		     * ever.
		     */
		    for(j = 0, t0 = strlen(tmfmt)*8; j < 3; j++, t0*=2) {
			addbufspc(t0);
			if ((len = ztrftime(bv->bp, t0, tmfmt, tm, ts.tv_nsec))
			    >= 0)
			    break;
		    }
		    /* There is enough room for this because addbufspc(t0)
		     * allocates room for t0 * 2 bytes. */
		    if (len >= 0)
			metafy(bv->bp, len, META_NOALLOC);
		    bv->bp += strlen(bv->bp);
		    zsfree(tmbuf);
		    break;
		}
	    case 'n':
		stradd(get_username());
		break;
	    case 'l':
		if (*ttystrname) {
                   ss = (strncmp(ttystrname, "/dev/tty", 8) ?
                           ttystrname + 5 : ttystrname + 8);
		    stradd(ss);
		} else
		    stradd("()");
		break;
	    case 'y':
		if (*ttystrname) {
		    ss = (strncmp(ttystrname, "/dev/", 5) ?
			    ttystrname : ttystrname + 5);
		    stradd(ss);
		} else
		    stradd("()");
		break;
	    case 'L':
		addbufspc(DIGBUFSIZE);
#if defined(ZLONG_IS_LONG_LONG) && defined(PRINTF_HAS_LLD)
		sprintf(bv->bp, "%lld", shlvl);
#else
		sprintf(bv->bp, "%ld", (long)shlvl);
#endif
		bv->bp += strlen(bv->bp);
		break;
	    case '?':
		addbufspc(DIGBUFSIZE);
#if defined(ZLONG_IS_LONG_LONG) && defined(PRINTF_HAS_LLD)
		sprintf(bv->bp, "%lld", lastval);
#else
		sprintf(bv->bp, "%ld", (long)lastval);
#endif
		bv->bp += strlen(bv->bp);
		break;
	    case '%':
	    case ')':
		addbufspc(1);
		*bv->bp++ = *bv->fm;
		break;
	    case '#':
		addbufspc(1);
		*bv->bp++ = privasserted() ? '#' : '%';
		break;
	    case 'v':
		if (!arg)
		    arg = 1;
		else if (arg < 0)
		    arg += arrlen(psvar) + 1;
		if (arg > 0 && arrlen_ge(psvar, arg))
		    stradd(psvar[arg - 1]);
		break;
	    case 'E':
                tsetcap(TCCLEAREOL, TSC_PROMPT);
		break;
	    case '^':
		if (cmdsp) {
		    if (arg >= 0) {
			if (arg > cmdsp || arg == 0)
			    arg = cmdsp;
			for (t0 = cmdsp - 1; arg--; t0--) {
			    stradd(cmdnames[cmdstack[t0]]);
			    if (arg) {
				addbufspc(1);
				*bv->bp++=' ';
			    }
			}
		    } else {
			arg = -arg;
			if (arg > cmdsp)
			    arg = cmdsp;
			for (t0 = arg - 1; arg--; t0--) {
			    stradd(cmdnames[cmdstack[t0]]);
			    if (arg) {
				addbufspc(1);
				*bv->bp++=' ';
			    }
			}
		    }
		}
		break;
	    case '_':
		if (cmdsp) {
		    if (arg >= 0) {
			if (arg > cmdsp || arg == 0)
			    arg = cmdsp;
			for (t0 = cmdsp - arg; arg--; t0++) {
			    stradd(cmdnames[cmdstack[t0]]);
			    if (arg) {
				addbufspc(1);
				*bv->bp++=' ';
			    }
			}
		    } else {
			arg = -arg;
			if (arg > cmdsp)
			    arg = cmdsp;
			for (t0 = 0; arg--; t0++) {
			    stradd(cmdnames[cmdstack[t0]]);
			    if (arg) {
				addbufspc(1);
				*bv->bp++=' ';
			    }
			}
		    }
		}
		break;
	    case 'r':
		if(bv->rstring)
		    stradd(bv->rstring);
		break;
	    case 'R':
		if(bv->Rstring)
		    stradd(bv->Rstring);
		break;
	    case 'e':
	    {
		int depth = 0;
		Funcstack fsptr = funcstack;
		while (fsptr) {
		    depth++;
		    fsptr = fsptr->prev;
		}
		addbufspc(DIGBUFSIZE);
		sprintf(bv->bp, "%d", depth);
		bv->bp += strlen(bv->bp);
		break;
	    }
	    case 'I':
		if (funcstack && funcstack->tp != FS_SOURCE &&
		    !IN_EVAL_TRAP()) {
		    /*
		     * We're in a function or an eval with
		     * EVALLINENO.  Calculate the line number in
		     * the file.
		     */
		    zlong flineno = lineno + funcstack->flineno;
		    /* take account of eval line nos. starting at 1 */
		    if (funcstack->tp == FS_EVAL)
			lineno--;
		    addbufspc(DIGBUFSIZE);
#if defined(ZLONG_IS_LONG_LONG) && defined(PRINTF_HAS_LLD)
		    sprintf(bv->bp, "%lld", flineno);
#else
		    sprintf(bv->bp, "%ld", (long)flineno);
#endif
		    bv->bp += strlen(bv->bp);
		    break;
		}
		/* else we're in a file and lineno is already correct */
		/* FALLTHROUGH */
	    case 'i':
		addbufspc(DIGBUFSIZE);
#if defined(ZLONG_IS_LONG_LONG) && defined(PRINTF_HAS_LLD)
		sprintf(bv->bp, "%lld", lineno);
#else
		sprintf(bv->bp, "%ld", (long)lineno);
#endif
		bv->bp += strlen(bv->bp);
		break;
	    case 'x':
		if (funcstack && funcstack->tp != FS_SOURCE &&
		    !IN_EVAL_TRAP())
		    promptpath(funcstack->filename ? funcstack->filename : "",
			       arg, 0);
		else
		    promptpath(scriptfilename ? scriptfilename : argzero,
			       arg, 0);
		break;
	    case '\0':
		return 0;
	    case Meta:
		bv->fm++;
		break;
	    }
	} else if(*bv->fm == '!' && isset(PROMPTBANG)) {
	    if(doprint) {
		if(bv->fm[1] == '!') {
		    bv->fm++;
		    addbufspc(1);
		    pputc('!');
		} else {
		    addbufspc(DIGBUFSIZE);
		    convbase(bv->bp, curhist, 10);
		    bv->bp += strlen(bv->bp);
		}
	    }
	} else {
	    char c = *bv->fm == Meta ? *++bv->fm ^ 32 : *bv->fm;

	    if (doprint) {
		addbufspc(1);
		pputc(c);
	    }
	}
    }

    return *bv->fm;
}


addbufspc(int need)
{
    need *= 2;   /* for metafication */
    if((bv->bp - bv->buf) + need > bv->bufspc) {
	int bo = bv->bp - bv->buf;
	int bo1 = bv->bp1 ? bv->bp1 - bv->buf : -1;
	ptrdiff_t bufline_off = bv->bufline ? bv->bufline - bv->buf : -1;

	if(need & 255)
	    need = (need | 255) + 1;
	bv->buf = realloc(bv->buf, bv->bufspc += need);
	memset(bv->buf + bv->bufspc - need, 0, need);
	bv->bp = bv->buf + bo;
	if(bo1 != -1)
	    bv->bp1 = bv->buf + bo1;
	if (bufline_off != -1)
	    bv->bufline = bv->buf + bufline_off;
    }
}


tsetcap(int cap, int flags)
{
    if (tccan(cap) && !isset(SINGLELINEZLE) &&
        !(termflags & (TERM_NOUP|TERM_BAD|TERM_UNKNOWN))) {
	switch (flags & TSC_OUTPUT_MASK) {
	case TSC_RAW:
	    tputs(tcstr[cap], 1, putraw);
	    break;
	case 0:
	default:
	    tputs(tcstr[cap], 1, putshout);
	    break;
	case TSC_PROMPT:
	    if (!bv->dontcount) {
		addbufspc(1);
		*bv->bp++ = Inpar;
	    }
	    tputs(tcstr[cap], 1, putstr);
	    if (!bv->dontcount) {
		int glitch = 0;

		if (cap == TCSTANDOUTBEG || cap == TCSTANDOUTEND)
		    glitch = tgetnum("sg");
		else if (cap == TCUNDERLINEBEG || cap == TCUNDERLINEEND)
		    glitch = tgetnum("ug");
		if(glitch < 0)
		    glitch = 0;
		addbufspc(glitch + 1);
		while(glitch--)
		    *bv->bp++ = Nularg;
		*bv->bp++ = Outpar;
	    }
	    break;
	}

	if (flags & TSC_DIRTY) {
	    flags &= ~TSC_DIRTY;
	    if (txtisset(TXTBOLDFACE) && cap != TCBOLDFACEBEG)
		tsetcap(TCBOLDFACEBEG, flags);
	    if (txtisset(TXTSTANDOUT))
		tsetcap(TCSTANDOUTBEG, flags);
	    if (txtisset(TXTUNDERLINE))
		tsetcap(TCUNDERLINEBEG, flags);
	    if (txtisset(TXTFGCOLOUR))
		set_colour_attribute(txtattrmask, COL_SEQ_FG, TSC_PROMPT);
	    if (txtisset(TXTBGCOLOUR))
		set_colour_attribute(txtattrmask, COL_SEQ_BG, TSC_PROMPT);
	}
    }
}


cmdpop(void)
{
    if (cmdsp <= 0) {
	DPUTS(1, "BUG: cmdstack empty");
	fflush(stderr);
    } else
	cmdsp--;
}


prompttrunc(int arg, int truncchar, int doprint, int endchar,
	    zattr *txtchangep)
{
    if (arg > 0) {
	char ch = *bv->fm, *ptr, *truncstr;
	int truncatleft = ch == '<';
	int w = bv->bp - bv->buf;

	/*
	 * If there is already a truncation active, return so that
	 * can be finished, backing up so that the new truncation
	 * can be started afterwards.
	 */
	if (bv->truncwidth) {
	    while (*--bv->fm != '%')
		;
	    bv->fm--;
	    return 0;
	}

	bv->truncwidth = arg;
	if (*bv->fm != ']')
	    bv->fm++;
	while (*bv->fm && *bv->fm != truncchar) {
	    if (*bv->fm == '\\' && bv->fm[1])
		++bv->fm;
	    addbufspc(1);
	    *bv->bp++ = *bv->fm++;
	}
	if (!*bv->fm)
	    return 0;
	if (bv->bp - bv->buf == w && truncchar == ']') {
	    addbufspc(1);
	    *bv->bp++ = '<';
	}
	ptr = bv->buf + w;	/* addbufspc() may have realloc()'d bv->buf */
	/*
	 * Now:
	 *   bv->buf is the start of the output prompt buffer
	 *   ptr is the start of the truncation string
	 *   bv->bp is the end of the truncation string
	 */
	truncstr = ztrduppfx(ptr, bv->bp - ptr);

	bv->bp = ptr;
	w = bv->bp - bv->buf;
	bv->fm++;
	bv->trunccount = bv->dontcount;
	putpromptchar(doprint, endchar, txtchangep);
	bv->trunccount = 0;
	ptr = bv->buf + w;	/* putpromptchar() may have realloc()'d */
	*bv->bp = '\0';
	/*
	 * Now:
	 *   ptr is the start of the truncation string and also
	 *     where we need to start putting any truncated output
	 *   bv->bp is the end of the string we have just added, which
	 *     may need truncating.
	 */

	/*
	 * w below is screen width if multibyte support is enabled
	 * (note that above it was a raw string pointer difference).
	 * It's the full width of the string we may need to truncate.
	 *
	 * bv->truncwidth has come from the user, so we interpret this
	 * as a screen width, too.
	 */
	countprompt(ptr, &w, 0, -1);
	if (w > bv->truncwidth) {
	    /*
	     * We need to truncate.  t points to the truncation string
	     * -- which is inserted literally, without nice
	     * representation.  twidth is its printing width, and maxwidth
	     * is the amount of the main string that we want to keep.
	     * Note that if the truncation string is longer than the
	     * truncation length (twidth > bv->truncwidth), the truncation
	     * string is used in full.
	     */
	    char *t = truncstr;
	    int fullen = bv->bp - ptr;
	    int twidth, maxwidth;
	    int ntrunc = strlen(t);

	    twidth = MB_METASTRWIDTH(t);
	    if (twidth < bv->truncwidth) {
		maxwidth = bv->truncwidth - twidth;
		/*
		 * It's not safe to assume there are no invisible substrings
		 * just because the width is less than the full string
		 * length since there may be multibyte characters.
		 */
		addbufspc(ntrunc+1);
		/* may have realloc'd */
		ptr = bv->bp - fullen;

		if (truncatleft) {
		    /*
		     * To truncate at the left, selectively copy
		     * maxwidth bytes from the main prompt, preceded
		     * by the truncation string in full.
		     *
		     * We're overwriting the string containing the
		     * text to be truncated, so copy it.  We've
		     * just ensured there's sufficient space at the
		     * end of the prompt string.
		     *
		     * Pointer into text to be truncated.
		     */
		    char *fulltextptr, *fulltext;
		    int remw;
#ifdef MULTIBYTE_SUPPORT
		    mbstate_t mbs;
		    memset(&mbs, 0, sizeof mbs);
#endif

		    fulltextptr = fulltext = ptr + ntrunc;
		    memmove(fulltext, ptr, fullen);
		    fulltext[fullen] = '\0';

		    /* Copy the truncstr into place. */
		    while (*t)
			*ptr++ = *t++;

		    /*
		     * Find the point in the text at which we should
		     * start copying, i.e. when the remaining width
		     * is less than or equal to the maximum width.
		     */
		    remw = w;
		    while (remw > maxwidth && *fulltextptr) {
			if (*fulltextptr == Inpar) {
			    /*
			     * Text marked as invisible: copy
			     * regardless, since we don't know what
			     * this does.  It only affects the width
			     * if there are Nularg's present.
			     * However, even in that case we
			     * can't break the sequence down, so
			     * we still loop over the entire group.
			     */
			    for (;;) {
				*ptr++ = *fulltextptr;
				if (*fulltextptr == '\0' ||
				    *fulltextptr++ == Outpar)
				    break;
				if (fulltextptr[-1] == Nularg)
				    remw--;
			    }
			} else {
#ifdef MULTIBYTE_SUPPORT
			    /*
			     * Normal text: build up a multibyte character.
			     */
			    char inchar;
			    wchar_t cc;
			    int wcw;

			    /*
			     * careful: string is still metafied (we
			     * need that because we don't know a
			     * priori when to stop and the resulting
			     * string must be metafied).
			     */
			    if (*fulltextptr == Meta)
				inchar = *++fulltextptr ^ 32;
			    else
				inchar = *fulltextptr;
			    fulltextptr++;
			    switch (mbrtowc(&cc, &inchar, 1, &mbs)) {
			    case MB_INCOMPLETE:
				/* Incomplete multibyte character. */
				break;
			    case MB_INVALID:
				/* Reset invalid state. */
				memset(&mbs, 0, sizeof mbs);
				/* FALL THROUGH */
			    case 0:
				/* Assume a single-byte character. */
				remw--;
				break;
			    default:
				wcw = WCWIDTH(cc);
				if (wcw >= 0)
				    remw -= wcw;
				else
				    remw--;
				break;
			    }
#else
			    /* Single byte character */
			    if (*fulltextptr == Meta)
				fulltextptr++;
			    fulltextptr++;
			    remw--;
#endif
			}
		    }

		    /*
		     * Now simply copy the rest of the text.  Still
		     * metafied, so this is easy.
		     */
		    while (*fulltextptr)
			*ptr++ = *fulltextptr++;
		    /* Mark the end of copying */
		    bv->bp = ptr;
		} else {
		    /*
		     * Truncating at the right is easier: just leave
		     * enough characters until we have reached the
		     * maximum width.
		     */
		    char *skiptext = ptr;
#ifdef MULTIBYTE_SUPPORT
		    mbstate_t mbs;
		    memset(&mbs, 0, sizeof mbs);
#endif

		    while (maxwidth > 0 && *skiptext) {
			if (*skiptext == Inpar) {
			    /* see comment on left truncation above */
			    for (;;) {
				if (*skiptext == '\0' ||
				    *skiptext++ == Outpar)
				    break;
				if (skiptext[-1] == Nularg)
				    maxwidth--;
			    }
			} else {
#ifdef MULTIBYTE_SUPPORT
			    char inchar;
			    wchar_t cc;
			    int wcw;

			    if (*skiptext == Meta)
				inchar = *++skiptext ^ 32;
			    else
				inchar = *skiptext;
			    skiptext++;
			    switch (mbrtowc(&cc, &inchar, 1, &mbs)) {
			    case MB_INCOMPLETE:
				/* Incomplete character. */
				break;
			    case MB_INVALID:
				/* Reset invalid state. */
				memset(&mbs, 0, sizeof mbs);
				/* FALL THROUGH */
			    case 0:
				/* Assume a single-byte character. */
				maxwidth--;
				break;
			    default:
				wcw = WCWIDTH(cc);
				if (wcw >= 0)
				    maxwidth -= wcw;
				else
				    maxwidth--;
				break;
			    }
#else
			    if (*skiptext == Meta)
				skiptext++;
			    skiptext++;
			    maxwidth--;
#endif
			}
		    }
		    /*
		     * We don't need the visible text from now on,
		     * but we'd better copy any invisible bits.
		     * History dictates that these go after the
		     * truncation string.  This is sensible since
		     * they may, for example, turn off an effect which
		     * should apply to all text at this point.
		     *
		     * Copy the truncstr.
		     */
		    ptr = skiptext;
		    while (*t)
			*ptr++ = *t++;
		    bv->bp = ptr;
		    if (*skiptext) {
			/* Move remaining text so we don't overwrite it */
			memmove(bv->bp, skiptext, strlen(skiptext)+1);
			skiptext = bv->bp;

			/*
			 * Copy anything we want, updating bv->bp
			 */
			while (*skiptext) {
			    if (*skiptext == Inpar) {
				for (;;) {
				    *bv->bp++ = *skiptext;
				    if (*skiptext == Outpar ||
					*skiptext == '\0')
					break;
				    skiptext++;
				}
			    }
			    else
				skiptext++;
			}
		    }
		}
	    } else {
		/* Just copy truncstr; no other text appears. */
		while (*t)
		    *ptr++ = *t++;
		bv->bp = ptr;
	    }
	    *bv->bp = '\0';
	}
	zsfree(truncstr);
	bv->truncwidth = 0;
	/*
	 * We may have returned early from the previous putpromptchar *
	 * because we found another truncation following this one.    *
	 * In that case we need to do the rest now.                   *
	 */
	if (!*bv->fm)
	    return 0;
	if (*bv->fm != endchar) {
	    bv->fm++;
	    /*
	     * With bv->truncwidth set to zero, we always reach endchar *
	     * (or the terminating NULL) this time round.         *
	     */
	    if (!putpromptchar(doprint, endchar, txtchangep))
		return 0;
	}
	/* Now we have to trick it into matching endchar again */
	bv->fm--;
    } else {
	if (*bv->fm != endchar)
	    bv->fm++;
	while(*bv->fm && *bv->fm != truncchar) {
	    if (*bv->fm == '\\' && bv->fm[1])
		bv->fm++;
	    bv->fm++;
	}
	if (bv->truncwidth || !*bv->fm)
	    return 0;
    }
    return 1;
}


promptexpand(char *s, int ns, char *rs, char *Rs, zattr *txtchangep)
{
    struct buf_vars new_vars;

    if(!s)
	return ztrdup("");

    if ((termflags & TERM_UNKNOWN) && (unset(INTERACTIVE)))
        init_term();

    if (isset(PROMPTSUBST)) {
	int olderr = errflag;
	int oldval = lastval;

	s = dupstring(s);
	if (!parsestr(&s))
	    singsub(&s);
	/*
	 * We don't need the special Nularg hack here and we're
	 * going to be using Nularg for other things.
	 */
	if (*s == Nularg && s[1] == '\0')
	    *s = '\0';

	/*
	 * Ignore errors and status change in prompt substitution.
	 * However, keep any user interrupt error that occurred.
	 */
	errflag = olderr | (errflag & ERRFLAG_INT);
	lastval = oldval;
    }

    memset(&new_vars, 0, sizeof(new_vars));
    new_vars.last = bv;
    bv = &new_vars;

    new_vars.rstring = rs;
    new_vars.Rstring = Rs;
    new_vars.fm = s;
    new_vars.bufspc = 256;
    new_vars.bp = new_vars.bufline = new_vars.buf = zshcalloc(new_vars.bufspc);
    new_vars.bp1 = NULL;
    new_vars.truncwidth = 0;

    putpromptchar(1, '\0', txtchangep);
    addbufspc(2);
    if (new_vars.dontcount)
	*new_vars.bp++ = Outpar;
    *new_vars.bp = '\0';
    if (!ns) {
	/* If zero, Inpar, Outpar and Nularg should be removed. */
	for (new_vars.bp = new_vars.buf; *new_vars.bp; ) {
	    if (*new_vars.bp == Meta)
		new_vars.bp += 2;
	    else if (*new_vars.bp == Inpar || *new_vars.bp == Outpar ||
		     *new_vars.bp == Nularg)
		chuck(new_vars.bp);
	    else
		new_vars.bp++;
	}
    }

    bv = new_vars.last;

    return new_vars.buf;
}


set_colour_attribute(zattr atr, int fg_bg, int flags)
{
    char *ptr;
    int do_free, is_prompt = (flags & TSC_PROMPT) ? 1 : 0;
    int colour, tc, def, use_termcap, use_truecolor;
    int is_default_zle_highlight = 1;

    if (fg_bg == COL_SEQ_FG) {
	colour = txtchangeget(atr, TXT_ATTR_FG_COL);
	tc = TCFGCOLOUR;
	def = txtchangeisset(atr, TXTNOFGCOLOUR);
	use_truecolor = txtchangeisset(atr, TXT_ATTR_FG_24BIT);
	use_termcap = txtchangeisset(atr, TXT_ATTR_FG_TERMCAP);
    } else {
	colour = txtchangeget(atr, TXT_ATTR_BG_COL);
	tc = TCBGCOLOUR;
	def = txtchangeisset(atr, TXTNOBGCOLOUR);
	use_truecolor = txtchangeisset(atr, TXT_ATTR_BG_24BIT);
	use_termcap = txtchangeisset(atr, TXT_ATTR_BG_TERMCAP);
    }

    /* Test if current zle_highlight settings are customized, or
     * the typical "standard" codes */
    if (0 != strcmp(fg_bg_sequences[fg_bg].start, fg_bg == COL_SEQ_FG ? TC_COL_FG_START : TC_COL_BG_START) ||
	/* the same in-fix for both FG and BG */
	0 != strcmp(fg_bg_sequences[fg_bg].def, TC_COL_FG_DEFAULT) ||
	/* the same suffix for both FG and BG */
	0 != strcmp(fg_bg_sequences[fg_bg].end, TC_COL_FG_END))
    {
	is_default_zle_highlight = 0;
    }

    /*
     * If we're not restoring the default, and either have a
     * colour value that is too large for ANSI, or have been told
     * to use the termcap sequence, try to use the termcap sequence.
     * True color is not covered by termcap.
     *
     * We have already sanitised the values we allow from the
     * highlighting variables, so much of this shouldn't be
     * necessary at this point, but we might as well be safe.
     */
    if (!def && !use_truecolor &&
	(is_default_zle_highlight && (colour > 7 || use_termcap)))
    {
	/*
	 * We can if it's available, and either we couldn't get
	 * the maximum number of colours, or the colour is in range.
	 */
	if (tccan(tc) && (tccolours < 0 || colour < tccolours))
	{
	    if (is_prompt)
	    {
		if (!bv->dontcount) {
		    addbufspc(1);
		    *bv->bp++ = Inpar;
		}
		tputs(tgoto(tcstr[tc], colour, colour), 1, putstr);
		if (!bv->dontcount) {
		    addbufspc(1);
		    *bv->bp++ = Outpar;
		}
	    } else {
		tputs(tgoto(tcstr[tc], colour, colour), 1, putshout);
	    }
	    /* That worked. */
	    return;
	}
	/*
	 * Nope, that didn't work.
	 * If 0 to 7, assume standard ANSI works, if 8 to 255, assume
         * typical 256-color escapes works, otherwise it won't.
	 */
	if (colour > 255)
	    return;
    }

    if ((do_free = (colseq_buf == NULL))) {
	/* This can happen when moving the cursor in trashzle() */
	allocate_colour_buffer();
    }

    /* Build the reset-code: .start + .def + . end
     * or the typical true-color code: .start + 8;2;%d;%d;%d + .end
     * or the typical 256-color code: .start + 8;5;%d + .end
     */
    if (use_truecolor)
	strcpy(colseq_buf, fg_bg == COL_SEQ_FG ? TC_COL_FG_START : TC_COL_BG_START);
    else
	strcpy(colseq_buf, fg_bg_sequences[fg_bg].start);

    ptr = colseq_buf + strlen(colseq_buf);
    if (def) {
	if (use_truecolor)
	    strcpy(ptr, fg_bg == COL_SEQ_FG ? TC_COL_FG_DEFAULT : TC_COL_BG_DEFAULT);
	else
	    strcpy(ptr, fg_bg_sequences[fg_bg].def);
	while (*ptr)
	    ptr++;
    } else if (use_truecolor) {
	ptr += sprintf(ptr, "8;2;%d;%d;%d", colour >> 16,
		(colour >> 8) & 0xff, colour & 0xff);
    } else if (colour > 7 && colour <= 255) {
	ptr += sprintf(ptr, "%d", colour);
    } else
	*ptr++ = colour + '0';
    if (use_truecolor)
	strcpy(ptr, fg_bg == COL_SEQ_FG ? TC_COL_FG_END : TC_COL_BG_END);
    else
	strcpy(ptr, fg_bg_sequences[fg_bg].end);

    if (is_prompt) {
	if (!bv->dontcount) {
	    addbufspc(1);
	    *bv->bp++ = Inpar;
	}
	tputs(colseq_buf, 1, putstr);
	if (!bv->dontcount) {
	    addbufspc(1);
	    *bv->bp++ = Outpar;
	}
    } else
	tputs(colseq_buf, 1, putshout);

    if (do_free)
	free_colour_buffer();
}


cmdpush(int cmdtok)
{
    if (cmdsp >= 0 && cmdsp < CMDSTACKSZ)
	cmdstack[cmdsp++] = (unsigned char)cmdtok;
}


static std::string getDescription(const std::string &fulltext, int line)
{
	if (line < 1) return "";

	unsigned int start = 0;
	for (; start<fulltext.length() ; ++start) {
		if (line <= 1) break;
		if (fulltext[start] == '\n') line--;
	}

	// not a valid description
	if (fulltext.compare(start, 2, "//") != 0) return "";

	// Jump over the two forward slashes
	start = start+2;

	//Jump over all the spaces
	while (fulltext[start] == ' ' || fulltext[start] == '\t') start++;
	std::string retString = "";

	// go till the end of the line
	while (fulltext[start] != '\n') {
		// replace // with space
		if (fulltext.compare(start, 2, "//") == 0) {
			retString += " ";
			start++;
		} else {
			retString += fulltext[start];
		}
		start++;
	}
	return retString;
}

Add file bounds check to comment parser


static GroupList collectGroups(const std::string &fulltext)
{
	GroupList groupList; // container of all group names
	int lineNo = 1; // tracks line number
	bool inString = false; // check if its string or (line-) comment

	// iterate through whole scad file
	for (unsigned int i=0; i<fulltext.length(); ++i) {
		// increase line number
		if (fulltext[i] == '\n') {
			lineNo++;
			continue;
		}

		// skip escaped quotes inside strings
		if (inString && fulltext.compare(i, 2, "\\\"") == 0) {
			i++;
			continue;
		}

		//start or end of string negate the checkpoint
		if (fulltext[i] == '"') {
			inString = !inString;
			continue;
		}

		if (!inString && fulltext.compare(i, 2, "//") == 0) {
			i++;
			while (fulltext[i] != '\n' && i<fulltext.length() ) i++;
			lineNo++;
			continue;
		}

		//start of multi line comment if check is true
		if (!inString && fulltext.compare(i, 2, "/*") == 0) {
			//store comment
			std::string comment;
			i++;
			if(i<fulltext.length()) {
				i++;
			}
			else {
				continue;
			}
			bool isGroup=true;
			// till */ every character is comment
			while (fulltext.compare(i, 2, "*/") != 0 && i<fulltext.length()) {
				if(fulltext[i]=='\n'){
					lineNo++;
					isGroup=false;
				}
				comment += fulltext[i];
				i++;
			}

			if(isGroup)
				groupList.push_back(createGroup(comment,lineNo));
		}
	}
	return groupList;
}

Add file bounds check to comment parser


void CommentParser::collectParameters(const std::string& fulltext, SourceFile *root_file)
{
	static auto EmptyStringLiteral(std::make_shared<Literal>(""));

	// Get all groups of parameters
	GroupList groupList = collectGroups(fulltext);
	int parseTill=getLineToStop(fulltext);
	// Extract parameters for all literal assignments
	for (auto &assignment : root_file->scope.assignments) {
		if (!assignment->getExpr()->isLiteral()) continue; // Only consider literals

		// get location of assignment node
		int firstLine = assignment->location().firstLine();
		if(firstLine>=parseTill || (
			assignment->location().fileName() != "" &&
			assignment->location().fileName() != root_file->getFilename() &&
			assignment->location().fileName() != root_file->getFullpath()
			)) {
			continue;
		}
		// making list to add annotations
		AnnotationList *annotationList = new AnnotationList();

		// Extracting the parameter comment
		std::string comment = getComment(fulltext, firstLine);
		// getting the node for parameter annotation
		shared_ptr<Expression> params = CommentParser::parser(comment.c_str());
		if (!params) {
			params = EmptyStringLiteral;
		}

		// adding parameter to the list
		annotationList->push_back(Annotation("Parameter", params));

		//extracting the description
		std::string descr = getDescription(fulltext, firstLine - 1);
		if (descr != "") {
			//creating node for description
			shared_ptr<Expression> expr(new Literal(descr));
			annotationList->push_back(Annotation("Description", expr));
		}

		// Look for the group to which the given assignment belong
		for (const auto &groupInfo :boost::adaptors::reverse(groupList)){
			if (groupInfo.lineNo < firstLine) {
				//creating node for description
				shared_ptr<Expression> expr(new Literal(groupInfo.commentString));
				annotationList->push_back(Annotation("Group", expr));
				break;
			}
		}
		assignment->addAnnotations(annotationList);
	}
}

Add file bounds check to comment parser


static int getLineToStop( const std::string &fulltext){
	int lineNo=1;
	bool inString=false;
	for (unsigned int i=0; i<fulltext.length(); ++i) {
	// increase line number
		if (fulltext[i] == '\n') {
			lineNo++;
			continue;
		}

		// skip escaped quotes inside strings
		if (inString && fulltext.compare(i, 2, "\\\"") == 0) {
			i++;
			continue;
		}

		//start or end of string negate the checkpoint
		if (fulltext[i] == '"') {
			inString = !inString;
			continue;
		}

		if (!inString && fulltext.compare(i, 2, "//") == 0) {
			i++;
			while (fulltext[i] != '\n' && i<fulltext.length()) i++;
			lineNo++;
			continue;
		}

		//start of multi line comment if check is true
		if (!inString && fulltext.compare(i, 2, "/*") == 0) {
			i ++;
			if(i<fulltext.length()) {
				i++;
			}
			else {
				continue;
			}
			// till */ every character is comment
			while (fulltext.compare(i, 2, "*/") != 0 && i<fulltext.length()) {
				if(fulltext[i]=='\n'){
					lineNo++;
				}
				i++;
			}
		}

		if (fulltext[i]== '{') {
			return lineNo;
		}
	}
	return lineNo;
}

Add file bounds check to comment parser


static GroupInfo createGroup(std::string comment,int lineNo)
{
	//store info related to group
	GroupInfo groupInfo;
	std::string finalGroupName;

	boost::regex regex("\\[(.*?)\\]");
	boost::match_results<std::string::const_iterator>  match;
	while(boost::regex_search(comment, match, regex)) {
		std::string groupName = match[1].str();
		if (finalGroupName.empty()) {
			finalGroupName = groupName;
		} else {
			finalGroupName = finalGroupName + "-" + groupName;
		}
		groupName.clear();
		comment = match.suffix();
	}

	groupInfo.commentString = finalGroupName;
	groupInfo.lineNo = lineNo;
	return groupInfo;
}

Add file bounds check to comment parser


static std::string getComment(const std::string &fulltext, int line)
{
	if (line < 1) return "";

	// Locate line
	unsigned int start = 0;
	for (; start<fulltext.length() ; ++start) {
		if (line <= 1) break;
		if (fulltext[start] == '\n') line--;
	}

	int end = start + 1;
	while (end < fulltext.size() && fulltext[end] != '\n') end++;

	std::string comment = fulltext.substr(start, end - start);

	// Locate comment
	unsigned int startText = 0;
	int noOfSemicolon = 0;
	bool inString = false;
	for (; startText < comment.length() - 1; ++startText) {
		if (inString && comment.compare(startText, 2, "\\\"") == 0) {
			startText++;
			continue;
		}
		if (comment[startText] == '"') inString = !inString;
		if (!inString) {
			if (comment.compare(startText, 2, "//") == 0) break;
			if (comment[startText] == ';' && noOfSemicolon > 0) return "";
			if (comment[startText] == ';') noOfSemicolon++;
		}
	}

	if (startText + 2 > comment.length()) return "";

	std::string result = comment.substr(startText + 2);
	return result;
}

Add file bounds check to comment parser


TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
  const auto* params =
      reinterpret_cast<const TfLiteGatherParams*>(node->builtin_data);
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
  const TfLiteTensor* positions;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputPositions, &positions));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensor, &output));

  if (positions->type == kTfLiteInt32) {
    switch (input->type) {
      case kTfLiteFloat32:
        return Gather<float, int32_t>(context, *params, input, positions,
                                      output);
      case kTfLiteUInt8:
        return Gather<uint8_t, int32_t>(context, *params, input, positions,
                                        output);
      case kTfLiteInt8:
        return Gather<int8_t, int32_t>(context, *params, input, positions,
                                       output);
      case kTfLiteInt16:
        return Gather<int16_t, int32_t>(context, *params, input, positions,
                                        output);
      case kTfLiteInt32:
        return Gather<int32_t, int32_t>(context, *params, input, positions,
                                        output);
      case kTfLiteInt64:
        return Gather<int64_t, int32_t>(context, *params, input, positions,
                                        output);
      case kTfLiteBool:
        return Gather<bool, int32_t>(context, *params, input, positions,
                                     output);
      case kTfLiteString:
        return GatherStrings<int32_t>(context, input, positions, output);
      default:
        context->ReportError(context, "Type '%s' is not supported by gather.",
                             TfLiteTypeGetName(input->type));
        return kTfLiteError;
    }
  }
  if (positions->type == kTfLiteInt64) {
    switch (input->type) {
      case kTfLiteFloat32:
        return Gather<float, int64_t>(context, *params, input, positions,
                                      output);
      case kTfLiteUInt8:
        return Gather<uint8_t, int64_t>(context, *params, input, positions,
                                        output);
      case kTfLiteInt8:
        return Gather<int8_t, int64_t>(context, *params, input, positions,
                                       output);
      case kTfLiteInt16:
        return Gather<int16_t, int64_t>(context, *params, input, positions,
                                        output);
      case kTfLiteInt32:
        return Gather<int32_t, int64_t>(context, *params, input, positions,
                                        output);
      case kTfLiteInt64:
        return Gather<int64_t, int64_t>(context, *params, input, positions,
                                        output);
      case kTfLiteBool:
        return Gather<bool, int64_t>(context, *params, input, positions,
                                     output);
      case kTfLiteString:
        return GatherStrings<int64_t>(context, input, positions, output);
      default:
        context->ReportError(context, "Type '%s' is not supported by gather.",
                             TfLiteTypeGetName(input->type));
        return kTfLiteError;
    }
  }
  context->ReportError(context,
                       "Positions of type '%s' are not supported by gather.",
                       TfLiteTypeGetName(positions->type));
  return kTfLiteError;
}

Prevent heap OOB read in TFLite's `gather.cc`.

Passing negative indices is illegal but there was a missing check so that resulted in OOB accesses.

PiperOrigin-RevId: 387231300
Change-Id: I3111b54b2f232638d795be17efc46abe4ede6bf8


TfLiteStatus GatherStrings(TfLiteContext* context, const TfLiteTensor* input,
                           const TfLiteTensor* positions,
                           TfLiteTensor* output) {
  DynamicBuffer buffer;

  const PositionT* indexes = GetTensorData<PositionT>(positions);
  bool indices_has_only_positive_elements = true;
  const size_t num_indices = positions->bytes / sizeof(PositionT);
  for (size_t i = 0; i < num_indices; i++) {
    if (indexes[i] < 0) {
      indices_has_only_positive_elements = false;
      break;
    }
  }
  TF_LITE_ENSURE(context, indices_has_only_positive_elements);

  const PositionT num_strings = GetStringCount(input);
  const int num_indexes = NumElements(positions);

  for (int i = 0; i < num_indexes; ++i) {
    const PositionT pos = indexes[i];
    TF_LITE_ENSURE(context, pos < num_strings);
    const auto string_ref = GetString(input, pos);
    buffer.AddString(string_ref.str, string_ref.len);
  }
  buffer.WriteToTensor(output, /*new_shape=*/nullptr);
  return kTfLiteOk;
}


TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);

  const auto* params =
      reinterpret_cast<const TfLiteGatherParams*>(node->builtin_data);
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
  const TfLiteTensor* positions;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputPositions, &positions));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputTensor, &output));

  switch (positions->type) {
    case kTfLiteInt64:
    case kTfLiteInt32:
      break;
    default:
      context->ReportError(
          context, "Positions of type '%s' are not supported by gather.",
          TfLiteTypeGetName(positions->type));
      return kTfLiteError;
  }

  // Assign to output the input type.
  output->type = input->type;

  // Check conditions for different types.
  switch (input->type) {
    case kTfLiteFloat32:
    case kTfLiteUInt8:
    case kTfLiteInt8:
    case kTfLiteInt16:
    case kTfLiteInt64:
    case kTfLiteInt32:
    case kTfLiteBool:
      break;
    case kTfLiteString: {
      // Only 1D input is supported.
      TF_LITE_ENSURE_EQ(context, NumDimensions(input), 1);
    } break;
    default:
      context->ReportError(context, "Type '%s' is not supported by gather.",
                           TfLiteTypeGetName(input->type));
      return kTfLiteError;
  }

  int axis = params->axis;
  if (axis < 0) {
    axis += NumDimensions(input);
  }
  TF_LITE_ENSURE(context, 0 <= axis && axis < NumDimensions(input));

  int batch_dims = params->batch_dims;
  // batch_dims should be in range: [-rank(positions), rank(positions)].
  // Negative batch_dims is added with rank of positions.
  if (batch_dims < 0) {
    batch_dims += NumDimensions(positions);
  }
  TF_LITE_ENSURE(context, batch_dims <= axis);
  TF_LITE_ENSURE(context, 0 <= batch_dims && batch_dims < NumDimensions(input));
  TF_LITE_ENSURE(context, batch_dims <= NumDimensions(positions));
  for (int i = 0; i < batch_dims; ++i) {
    TF_LITE_ENSURE_EQ(context, input->dims->data[i], positions->dims->data[i]);
  }

  const int num_dimensions =
      NumDimensions(input) + NumDimensions(positions) - 1 - batch_dims;
  TfLiteIntArray* output_shape = TfLiteIntArrayCreate(num_dimensions);
  int output_index = 0;
  for (int i = 0; i < axis; ++i) {
    output_shape->data[output_index++] = input->dims->data[i];
  }
  for (int i = batch_dims; i < positions->dims->size; ++i) {
    output_shape->data[output_index++] = positions->dims->data[i];
  }
  for (int i = axis + 1; i < input->dims->size; ++i) {
    output_shape->data[output_index++] = input->dims->data[i];
  }
  return context->ResizeTensor(context, output, output_shape);
}


TfLiteStatus Gather(TfLiteContext* context, const TfLiteGatherParams& params,
                    const TfLiteTensor* input, const TfLiteTensor* positions,
                    TfLiteTensor* output) {
  const PositionsT* indexes = GetTensorData<PositionsT>(positions);
  bool indices_has_only_positive_elements = true;
  const size_t num_indices = positions->bytes / sizeof(PositionsT);
  for (size_t i = 0; i < num_indices; i++) {
    if (indexes[i] < 0) {
      indices_has_only_positive_elements = false;
      break;
    }
  }
  TF_LITE_ENSURE(context, indices_has_only_positive_elements);

  tflite::GatherParams op_params;
  op_params.axis = params.axis;
  op_params.batch_dims = params.batch_dims;
  optimized_ops::Gather(op_params, GetTensorShape(input),
                        GetTensorData<InputT>(input), GetTensorShape(positions),
                        GetTensorData<PositionsT>(positions),
                        GetTensorShape(output), GetTensorData<InputT>(output));
  return kTfLiteOk;
}


static int sqfs_read_entry(struct squashfs_directory_entry **dest, void *src)
{
	struct squashfs_directory_entry *tmp;
	u16 sz;

	tmp = src;
	sz = get_unaligned_le16(src + sizeof(*tmp) - sizeof(u16));
	/*
	 * 'src' points to the begin of a directory entry, and 'sz' gets its
	 * 'name_size' member's value. name_size is actually the string
	 * length - 1, so adding 2 compensates this difference and adds space
	 * for the trailling null byte.
	 */
	*dest = malloc(sizeof(*tmp) + sz + 2);
	if (!*dest)
		return -ENOMEM;

	memcpy(*dest, src, sizeof(*tmp) + sz + 1);
	(*dest)->name[sz + 1] = '\0';

	return 0;
}

fs/squashfs: sqfs_read: Prevent arbitrary code execution

Following Jincheng's report, an out-of-band write leading to arbitrary
code execution is possible because on one side the squashfs logic
accepts directory names up to 65535 bytes (u16), while U-Boot fs logic
accepts directory names up to 255 bytes long.

Prevent such an exploit from happening by capping directory name sizes
to 255. Use a define for this purpose so that developers can link the
limitation to its source and eventually kill it some day by dynamically
allocating this array (if ever desired).

Link: https://lore.kernel.org/all/CALO=DHFB+yBoXxVr5KcsK0iFdg+e7ywko4-e+72kjbcS8JBfPw@mail.gmail.com
Reported-by: Jincheng Wang <jc.w4ng@gmail.com>
Signed-off-by: Miquel Raynal <miquel.raynal@bootlin.com>
Tested-by: Jincheng Wang <jc.w4ng@gmail.com>


void sqfs_close(void)
{
	sqfs_decompressor_cleanup(&ctxt);
	free(ctxt.sblk);
	ctxt.sblk = NULL;
	ctxt.cur_dev = NULL;
}


int sqfs_read(const char *filename, void *buf, loff_t offset, loff_t len,
	      loff_t *actread)
{
	char *dir = NULL, *fragment_block, *datablock = NULL;
	char *fragment = NULL, *file = NULL, *resolved, *data;
	u64 start, n_blks, table_size, data_offset, table_offset, sparse_size;
	int ret, j, i_number, datablk_count = 0;
	struct squashfs_super_block *sblk = ctxt.sblk;
	struct squashfs_fragment_block_entry frag_entry;
	struct squashfs_file_info finfo = {0};
	struct squashfs_symlink_inode *symlink;
	struct fs_dir_stream *dirsp = NULL;
	struct squashfs_dir_stream *dirs;
	struct squashfs_lreg_inode *lreg;
	struct squashfs_base_inode *base;
	struct squashfs_reg_inode *reg;
	unsigned long dest_len;
	struct fs_dirent *dent;
	unsigned char *ipos;

	*actread = 0;

	if (offset) {
		/*
		 * TODO: implement reading at an offset in file
		 */
		printf("Error: reading at a specific offset in a squashfs file is not supported yet.\n");
		return -EINVAL;
	}

	/*
	 * sqfs_opendir will uncompress inode and directory tables, and will
	 * return a pointer to the directory that contains the requested file.
	 */
	sqfs_split_path(&file, &dir, filename);
	ret = sqfs_opendir(dir, &dirsp);
	if (ret) {
		goto out;
	}

	dirs = (struct squashfs_dir_stream *)dirsp;

	/* For now, only regular files are able to be loaded */
	while (!sqfs_readdir(dirsp, &dent)) {
		ret = strcmp(dent->name, file);
		if (!ret)
			break;

		free(dirs->entry);
		dirs->entry = NULL;
	}

	if (ret) {
		printf("File not found.\n");
		*actread = 0;
		ret = -ENOENT;
		goto out;
	}

	i_number = dirs->dir_header->inode_number + dirs->entry->inode_offset;
	ipos = sqfs_find_inode(dirs->inode_table, i_number, sblk->inodes,
			       sblk->block_size);

	base = (struct squashfs_base_inode *)ipos;
	switch (get_unaligned_le16(&base->inode_type)) {
	case SQFS_REG_TYPE:
		reg = (struct squashfs_reg_inode *)ipos;
		datablk_count = sqfs_get_regfile_info(reg, &finfo, &frag_entry,
						      sblk->block_size);
		if (datablk_count < 0) {
			ret = -EINVAL;
			goto out;
		}

		memcpy(finfo.blk_sizes, ipos + sizeof(*reg),
		       datablk_count * sizeof(u32));
		break;
	case SQFS_LREG_TYPE:
		lreg = (struct squashfs_lreg_inode *)ipos;
		datablk_count = sqfs_get_lregfile_info(lreg, &finfo,
						       &frag_entry,
						       sblk->block_size);
		if (datablk_count < 0) {
			ret = -EINVAL;
			goto out;
		}

		memcpy(finfo.blk_sizes, ipos + sizeof(*lreg),
		       datablk_count * sizeof(u32));
		break;
	case SQFS_SYMLINK_TYPE:
	case SQFS_LSYMLINK_TYPE:
		symlink = (struct squashfs_symlink_inode *)ipos;
		resolved = sqfs_resolve_symlink(symlink, filename);
		ret = sqfs_read(resolved, buf, offset, len, actread);
		free(resolved);
		goto out;
	case SQFS_BLKDEV_TYPE:
	case SQFS_CHRDEV_TYPE:
	case SQFS_LBLKDEV_TYPE:
	case SQFS_LCHRDEV_TYPE:
	case SQFS_FIFO_TYPE:
	case SQFS_SOCKET_TYPE:
	case SQFS_LFIFO_TYPE:
	case SQFS_LSOCKET_TYPE:
	default:
		printf("Unsupported entry type\n");
		ret = -EINVAL;
		goto out;
	}

	/* If the user specifies a length, check its sanity */
	if (len) {
		if (len > finfo.size) {
			ret = -EINVAL;
			goto out;
		}

		finfo.size = len;
	} else {
		len = finfo.size;
	}

	if (datablk_count) {
		data_offset = finfo.start;
		datablock = malloc(get_unaligned_le32(&sblk->block_size));
		if (!datablock) {
			ret = -ENOMEM;
			goto out;
		}
	}

	for (j = 0; j < datablk_count; j++) {
		char *data_buffer;

		start = lldiv(data_offset, ctxt.cur_dev->blksz);
		table_size = SQFS_BLOCK_SIZE(finfo.blk_sizes[j]);
		table_offset = data_offset - (start * ctxt.cur_dev->blksz);
		n_blks = DIV_ROUND_UP(table_size + table_offset,
				      ctxt.cur_dev->blksz);

		/* Don't load any data for sparse blocks */
		if (finfo.blk_sizes[j] == 0) {
			n_blks = 0;
			table_offset = 0;
			data_buffer = NULL;
			data = NULL;
		} else {
			data_buffer = malloc_cache_aligned(n_blks * ctxt.cur_dev->blksz);

			if (!data_buffer) {
				ret = -ENOMEM;
				goto out;
			}

			ret = sqfs_disk_read(start, n_blks, data_buffer);
			if (ret < 0) {
				/*
				 * Possible causes: too many data blocks or too large
				 * SquashFS block size. Tip: re-compile the SquashFS
				 * image with mksquashfs's -b <block_size> option.
				 */
				printf("Error: too many data blocks to be read.\n");
				goto out;
			}

			data = data_buffer + table_offset;
		}

		/* Load the data */
		if (finfo.blk_sizes[j] == 0) {
			/* This is a sparse block */
			sparse_size = get_unaligned_le32(&sblk->block_size);
			if ((*actread + sparse_size) > len)
				sparse_size = len - *actread;
			memset(buf + *actread, 0, sparse_size);
			*actread += sparse_size;
		} else if (SQFS_COMPRESSED_BLOCK(finfo.blk_sizes[j])) {
			dest_len = get_unaligned_le32(&sblk->block_size);
			ret = sqfs_decompress(&ctxt, datablock, &dest_len,
					      data, table_size);
			if (ret)
				goto out;

			if ((*actread + dest_len) > len)
				dest_len = len - *actread;
			memcpy(buf + *actread, datablock, dest_len);
			*actread += dest_len;
		} else {
			if ((*actread + table_size) > len)
				table_size = len - *actread;
			memcpy(buf + *actread, data, table_size);
			*actread += table_size;
		}

		data_offset += table_size;
		free(data_buffer);
		if (*actread >= len)
			break;
	}

	/*
	 * There is no need to continue if the file is not fragmented.
	 */
	if (!finfo.frag) {
		ret = 0;
		goto out;
	}

	start = lldiv(frag_entry.start, ctxt.cur_dev->blksz);
	table_size = SQFS_BLOCK_SIZE(frag_entry.size);
	table_offset = frag_entry.start - (start * ctxt.cur_dev->blksz);
	n_blks = DIV_ROUND_UP(table_size + table_offset, ctxt.cur_dev->blksz);

	fragment = malloc_cache_aligned(n_blks * ctxt.cur_dev->blksz);

	if (!fragment) {
		ret = -ENOMEM;
		goto out;
	}

	ret = sqfs_disk_read(start, n_blks, fragment);
	if (ret < 0)
		goto out;

	/* File compressed and fragmented */
	if (finfo.frag && finfo.comp) {
		dest_len = get_unaligned_le32(&sblk->block_size);
		fragment_block = malloc(dest_len);
		if (!fragment_block) {
			ret = -ENOMEM;
			goto out;
		}

		ret = sqfs_decompress(&ctxt, fragment_block, &dest_len,
				      (void *)fragment  + table_offset,
				      frag_entry.size);
		if (ret) {
			free(fragment_block);
			goto out;
		}

		memcpy(buf + *actread, &fragment_block[finfo.offset], finfo.size - *actread);
		*actread = finfo.size;

		free(fragment_block);

	} else if (finfo.frag && !finfo.comp) {
		fragment_block = (void *)fragment + table_offset;

		memcpy(buf + *actread, &fragment_block[finfo.offset], finfo.size - *actread);
		*actread = finfo.size;
	}

out:
	free(fragment);
	free(datablock);
	free(file);
	free(dir);
	free(finfo.blk_sizes);
	sqfs_closedir(dirsp);

	return ret;
}


static int sqfs_split_path(char **file, char **dir, const char *path)
{
	char *dirc, *basec, *bname, *dname, *tmp_path;
	int ret = 0;

	*file = NULL;
	*dir = NULL;
	dirc = NULL;
	basec = NULL;
	bname = NULL;
	dname = NULL;
	tmp_path = NULL;

	/* check for first slash in path*/
	if (path[0] == '/') {
		tmp_path = strdup(path);
		if (!tmp_path) {
			ret = -ENOMEM;
			goto out;
		}
	} else {
		tmp_path = malloc(strlen(path) + 2);
		if (!tmp_path) {
			ret = -ENOMEM;
			goto out;
		}
		tmp_path[0] = '/';
		strcpy(tmp_path + 1, path);
	}

	/* String duplicates */
	dirc = strdup(tmp_path);
	if (!dirc) {
		ret = -ENOMEM;
		goto out;
	}

	basec = strdup(tmp_path);
	if (!basec) {
		ret = -ENOMEM;
		goto out;
	}

	dname = sqfs_dirname(dirc);
	bname = sqfs_basename(basec);

	*file = strdup(bname);

	if (!*file) {
		ret = -ENOMEM;
		goto out;
	}

	if (*dname == '\0') {
		*dir = malloc(2);
		if (!*dir) {
			ret = -ENOMEM;
			goto out;
		}

		(*dir)[0] = '/';
		(*dir)[1] = '\0';
	} else {
		*dir = strdup(dname);
		if (!*dir) {
			ret = -ENOMEM;
			goto out;
		}
	}

out:
	if (ret) {
		free(*file);
		free(*dir);
		*dir = NULL;
		*file = NULL;
	}
	free(basec);
	free(dirc);
	free(tmp_path);

	return ret;
}


int sqfs_exists(const char *filename)
{
	struct fs_dir_stream *dirsp = NULL;
	struct squashfs_dir_stream *dirs;
	char *dir, *file;
	struct fs_dirent *dent;
	int ret;

	sqfs_split_path(&file, &dir, filename);
	/*
	 * sqfs_opendir will uncompress inode and directory tables, and will
	 * return a pointer to the directory that contains the requested file.
	 */
	ret = sqfs_opendir(dir, &dirsp);
	if (ret) {
		ret = -EINVAL;
		goto free_strings;
	}

	dirs = (struct squashfs_dir_stream *)dirsp;

	while (!sqfs_readdir(dirsp, &dent)) {
		ret = strcmp(dent->name, file);
		if (!ret)
			break;
		free(dirs->entry);
		dirs->entry = NULL;
	}

	sqfs_closedir(dirsp);

free_strings:
	free(dir);
	free(file);

	return ret == 0;
}


static int sqfs_calc_n_blks(__le64 start, __le64 end, u64 *offset)
{
	u64 start_, table_size;

	table_size = le64_to_cpu(end) - le64_to_cpu(start);
	start_ = le64_to_cpu(start) / ctxt.cur_dev->blksz;
	*offset = le64_to_cpu(start) - (start_ * ctxt.cur_dev->blksz);

	return DIV_ROUND_UP(table_size + *offset, ctxt.cur_dev->blksz);
}


int sqfs_size(const char *filename, loff_t *size)
{
	struct squashfs_super_block *sblk = ctxt.sblk;
	struct squashfs_symlink_inode *symlink;
	struct fs_dir_stream *dirsp = NULL;
	struct squashfs_base_inode *base;
	struct squashfs_dir_stream *dirs;
	struct squashfs_lreg_inode *lreg;
	struct squashfs_reg_inode *reg;
	char *dir, *file, *resolved;
	struct fs_dirent *dent;
	unsigned char *ipos;
	int ret, i_number;

	sqfs_split_path(&file, &dir, filename);
	/*
	 * sqfs_opendir will uncompress inode and directory tables, and will
	 * return a pointer to the directory that contains the requested file.
	 */
	ret = sqfs_opendir(dir, &dirsp);
	if (ret) {
		ret = -EINVAL;
		goto free_strings;
	}

	dirs = (struct squashfs_dir_stream *)dirsp;

	while (!sqfs_readdir(dirsp, &dent)) {
		ret = strcmp(dent->name, file);
		if (!ret)
			break;
		free(dirs->entry);
		dirs->entry = NULL;
	}

	if (ret) {
		printf("File not found.\n");
		*size = 0;
		ret = -EINVAL;
		goto free_strings;
	}

	i_number = dirs->dir_header->inode_number + dirs->entry->inode_offset;
	ipos = sqfs_find_inode(dirs->inode_table, i_number, sblk->inodes,
			       sblk->block_size);
	free(dirs->entry);
	dirs->entry = NULL;

	base = (struct squashfs_base_inode *)ipos;
	switch (get_unaligned_le16(&base->inode_type)) {
	case SQFS_REG_TYPE:
		reg = (struct squashfs_reg_inode *)ipos;
		*size = get_unaligned_le32(&reg->file_size);
		break;
	case SQFS_LREG_TYPE:
		lreg = (struct squashfs_lreg_inode *)ipos;
		*size = get_unaligned_le64(&lreg->file_size);
		break;
	case SQFS_SYMLINK_TYPE:
	case SQFS_LSYMLINK_TYPE:
		symlink = (struct squashfs_symlink_inode *)ipos;
		resolved = sqfs_resolve_symlink(symlink, filename);
		ret = sqfs_size(resolved, size);
		free(resolved);
		break;
	case SQFS_BLKDEV_TYPE:
	case SQFS_CHRDEV_TYPE:
	case SQFS_LBLKDEV_TYPE:
	case SQFS_LCHRDEV_TYPE:
	case SQFS_FIFO_TYPE:
	case SQFS_SOCKET_TYPE:
	case SQFS_LFIFO_TYPE:
	case SQFS_LSOCKET_TYPE:
	default:
		printf("Unable to recover entry's size.\n");
		*size = 0;
		ret = -EINVAL;
		break;
	}

free_strings:
	free(dir);
	free(file);

	sqfs_closedir(dirsp);

	return ret;
}


static int sqfs_get_tokens_length(char **tokens, int count)
{
	int length = 0, i;

	/*
	 * 1 is added to the result of strlen to consider the slash separator
	 * between the tokens.
	 */
	for (i = 0; i < count; i++)
		length += strlen(tokens[i]) + 1;

	return length;
}


static char *sqfs_dirname(char *path)
{
	char *fname;

	fname = sqfs_basename(path);
	--fname;
	*fname = '\0';

	return path;
}


static char *sqfs_resolve_symlink(struct squashfs_symlink_inode *sym,
				  const char *base_path)
{
	char *resolved, *target;
	u32 sz;

	sz = get_unaligned_le32(&sym->symlink_size);
	target = malloc(sz + 1);
	if (!target)
		return NULL;

	/*
	 * There is no trailling null byte in the symlink's target path, so a
	 * copy is made and a '\0' is added at its end.
	 */
	target[sz] = '\0';
	/* Get target name (relative path) */
	strncpy(target, sym->symlink, sz);

	/* Relative -> absolute path conversion */
	resolved = sqfs_get_abs_path(base_path, target);

	free(target);

	return resolved;
}


static char *sqfs_concat_tokens(char **token_list, int token_count)
{
	char *result;
	int i, length = 0, offset = 0;

	length = sqfs_get_tokens_length(token_list, token_count);

	result = malloc(length + 1);
	if (!result)
		return NULL;

	result[length] = '\0';

	for (i = 0; i < token_count; i++) {
		strcpy(result + offset, token_list[i]);
		offset += strlen(token_list[i]);
		result[offset++] = '/';
	}

	return result;
}


static int sqfs_read_inode_table(unsigned char **inode_table)
{
	struct squashfs_super_block *sblk = ctxt.sblk;
	u64 start, n_blks, table_offset, table_size;
	int j, ret = 0, metablks_count;
	unsigned char *src_table, *itb;
	u32 src_len, dest_offset = 0;
	unsigned long dest_len = 0;
	bool compressed;

	table_size = get_unaligned_le64(&sblk->directory_table_start) -
		get_unaligned_le64(&sblk->inode_table_start);
	start = get_unaligned_le64(&sblk->inode_table_start) /
		ctxt.cur_dev->blksz;
	n_blks = sqfs_calc_n_blks(sblk->inode_table_start,
				  sblk->directory_table_start, &table_offset);

	/* Allocate a proper sized buffer (itb) to store the inode table */
	itb = malloc_cache_aligned(n_blks * ctxt.cur_dev->blksz);
	if (!itb)
		return -ENOMEM;

	if (sqfs_disk_read(start, n_blks, itb) < 0) {
		ret = -EINVAL;
		goto free_itb;
	}

	/* Parse inode table (metadata block) header */
	ret = sqfs_read_metablock(itb, table_offset, &compressed, &src_len);
	if (ret) {
		ret = -EINVAL;
		goto free_itb;
	}

	/* Calculate size to store the whole decompressed table */
	metablks_count = sqfs_count_metablks(itb, table_offset, table_size);
	if (metablks_count < 1) {
		ret = -EINVAL;
		goto free_itb;
	}

	*inode_table = malloc(metablks_count * SQFS_METADATA_BLOCK_SIZE);
	if (!*inode_table) {
		ret = -ENOMEM;
		printf("Error: failed to allocate squashfs inode_table of size %i, increasing CONFIG_SYS_MALLOC_LEN could help\n",
		       metablks_count * SQFS_METADATA_BLOCK_SIZE);
		goto free_itb;
	}

	src_table = itb + table_offset + SQFS_HEADER_SIZE;

	/* Extract compressed Inode table */
	for (j = 0; j < metablks_count; j++) {
		sqfs_read_metablock(itb, table_offset, &compressed, &src_len);
		if (compressed) {
			dest_len = SQFS_METADATA_BLOCK_SIZE;
			ret = sqfs_decompress(&ctxt, *inode_table +
					      dest_offset, &dest_len,
					      src_table, src_len);
			if (ret) {
				free(*inode_table);
				*inode_table = NULL;
				goto free_itb;
			}

			dest_offset += dest_len;
		} else {
			memcpy(*inode_table + (j * SQFS_METADATA_BLOCK_SIZE),
			       src_table, src_len);
		}

		/*
		 * Offsets to the decompression destination, to the metadata
		 * buffer 'itb' and to the decompression source, respectively.
		 */

		table_offset += src_len + SQFS_HEADER_SIZE;
		src_table += src_len + SQFS_HEADER_SIZE;
	}

free_itb:
	free(itb);

	return ret;
}


int sqfs_opendir(const char *filename, struct fs_dir_stream **dirsp)
{
	unsigned char *inode_table = NULL, *dir_table = NULL;
	int j, token_count = 0, ret = 0, metablks_count;
	struct squashfs_dir_stream *dirs;
	char **token_list = NULL, *path = NULL;
	u32 *pos_list = NULL;

	dirs = calloc(1, sizeof(*dirs));
	if (!dirs)
		return -EINVAL;

	/* these should be set to NULL to prevent dangling pointers */
	dirs->dir_header = NULL;
	dirs->entry = NULL;
	dirs->table = NULL;
	dirs->inode_table = NULL;
	dirs->dir_table = NULL;

	ret = sqfs_read_inode_table(&inode_table);
	if (ret) {
		ret = -EINVAL;
		goto out;
	}

	metablks_count = sqfs_read_directory_table(&dir_table, &pos_list);
	if (metablks_count < 1) {
		ret = -EINVAL;
		goto out;
	}

	/* Tokenize filename */
	token_count = sqfs_count_tokens(filename);
	if (token_count < 0) {
		ret = -EINVAL;
		goto out;
	}

	path = strdup(filename);
	if (!path) {
		ret = -EINVAL;
		goto out;
	}

	token_list = malloc(token_count * sizeof(char *));
	if (!token_list) {
		ret = -EINVAL;
		goto out;
	}

	/* Fill tokens list */
	ret = sqfs_tokenize(token_list, token_count, path);
	if (ret)
		goto out;
	/*
	 * ldir's (extended directory) size is greater than dir, so it works as
	 * a general solution for the malloc size, since 'i' is a union.
	 */
	dirs->inode_table = inode_table;
	dirs->dir_table = dir_table;
	ret = sqfs_search_dir(dirs, token_list, token_count, pos_list,
			      metablks_count);
	if (ret)
		goto out;

	if (le16_to_cpu(dirs->i_dir.inode_type) == SQFS_DIR_TYPE)
		dirs->size = le16_to_cpu(dirs->i_dir.file_size);
	else
		dirs->size = le32_to_cpu(dirs->i_ldir.file_size);

	/* Setup directory header */
	memcpy(dirs->dir_header, dirs->table, SQFS_DIR_HEADER_SIZE);
	dirs->entry_count = dirs->dir_header->count + 1;
	dirs->size -= SQFS_DIR_HEADER_SIZE;

	/* Setup entry */
	dirs->entry = NULL;
	dirs->table += SQFS_DIR_HEADER_SIZE;

	*dirsp = (struct fs_dir_stream *)dirs;

out:
	for (j = 0; j < token_count; j++)
		free(token_list[j]);
	free(token_list);
	free(pos_list);
	free(path);
	if (ret) {
		free(inode_table);
		free(dirs);
	}

	return ret;
}